StorageTextureNode: Add TSL read/write support (#32734)

examples/webgpu_compute_texture_pingpong.html

@@ -16,7 +16,7 @@
 			</div>
 
 			<small>
-				Compute ping/pong texture using GPU.
+				Compute ping/pong texture using GPU (pure TSL).
 			</small>
 		</div>
 
@@ -34,7 +34,7 @@
 		<script type="module">
 
 			import * as THREE from 'three/webgpu';
-			import { storageTexture, wgslFn, code, instanceIndex, uniform, NodeAccess } from 'three/tsl';
+			import { storageTexture, textureStore, Fn, instanceIndex, uniform, float, vec2, vec4, uvec2, ivec2, int, NodeAccess } from 'three/tsl';
 
 			import WebGPU from 'three/addons/capabilities/WebGPU.js';
 
@@ -45,6 +45,8 @@
 			let phase = true;
 			let lastUpdate = - 1;
 
+			const width = 512, height = 512;
+
 			const seed = uniform( new THREE.Vector2() );
 
 			init();
@@ -68,7 +70,6 @@
 				// texture
 
 				const hdr = true;
-				const width = 512, height = 512;
 
 				pingTexture = new THREE.StorageTexture( width, height );
 				pongTexture = new THREE.StorageTexture( width, height );
@@ -80,83 +81,63 @@
 
 				}
 
-				const wgslFormat = hdr ? 'rgba16float' : 'rgba8unorm';
-			
-				const readPing = storageTexture( pingTexture ).setAccess( NodeAccess.READ_ONLY );
-				const writePing = storageTexture( pingTexture ).setAccess( NodeAccess.WRITE_ONLY );
-				const readPong = storageTexture( pongTexture ).setAccess( NodeAccess.READ_ONLY );
-				const writePong = storageTexture( pongTexture ).setAccess( NodeAccess.WRITE_ONLY );
-
-				// compute init
-
-				const rand2 = code( `
-					fn rand2( n: vec2f ) -> f32 {
-
-						return fract( sin( dot( n, vec2f( 12.9898, 4.1414 ) ) ) * 43758.5453 );
+				const rand2 = Fn( ( [ n ] ) => {
 
-					}
+					return n.dot( vec2( 12.9898, 4.1414 ) ).sin().mul( 43758.5453 ).fract();
 
-					fn blur( image : texture_storage_2d<${wgslFormat}, read>, uv : vec2i ) -> vec4f {
+				} );
 
-						var color = vec4f( 0.0 );
-
-						color += textureLoad( image, uv + vec2i( - 1, 1 ));
-						color += textureLoad( image, uv + vec2i( - 1, - 1 ));
-						color += textureLoad( image, uv + vec2i( 0, 0 ));
-						color += textureLoad( image, uv + vec2i( 1, - 1 ));
-						color += textureLoad( image, uv + vec2i( 1, 1 ));
-
-						return color / 5.0; 
-					}
-
-					fn getUV( posX: u32, posY: u32 ) -> vec2f {
-
-						let uv = vec2f( f32( posX ) / ${ width }.0, f32( posY ) / ${ height }.0 );
+				// Create storage texture nodes with proper access
+				const writePing = storageTexture( pingTexture ).setAccess( NodeAccess.WRITE_ONLY );
+				const readPing = storageTexture( pingTexture ).setAccess( NodeAccess.READ_ONLY );
+				const writePong = storageTexture( pongTexture ).setAccess( NodeAccess.WRITE_ONLY );
+				const readPong = storageTexture( pongTexture ).setAccess( NodeAccess.READ_ONLY );
 
-						return uv;
+				const computeInit = Fn( () => {
 
-					}
-				` );
+					const posX = instanceIndex.mod( width );
+					const posY = instanceIndex.div( width );
+					const indexUV = uvec2( posX, posY );
+					const uv = vec2( float( posX ).div( width ), float( posY ).div( height ) );
 
-				const computeInitWGSL = wgslFn( `
-					fn computeInitWGSL( writeTex: texture_storage_2d<${ wgslFormat }, write>, index: u32, seed: vec2f ) -> void {
+					const r = rand2( uv.add( seed.mul( 100 ) ) ).sub( rand2( uv.add( seed.mul( 300 ) ) ) );
+					const g = rand2( uv.add( seed.mul( 200 ) ) ).sub( rand2( uv.add( seed.mul( 300 ) ) ) );
+					const b = rand2( uv.add( seed.mul( 200 ) ) ).sub( rand2( uv.add( seed.mul( 100 ) ) ) );
 
-						let posX = index % ${ width };
-						let posY = index / ${ width };
-						let indexUV = vec2u( posX, posY );
-						let uv = getUV( posX, posY );
+					textureStore( writePing, indexUV, vec4( r, g, b, 1 ) );
 
-						let r = rand2( uv + seed * 100 ) - rand2( uv + seed * 300 );
-						let g = rand2( uv + seed * 200 ) - rand2( uv + seed * 300 );
-						let b = rand2( uv + seed * 200 ) - rand2( uv + seed * 100 );
+				} );
 
-						textureStore( writeTex, indexUV, vec4( r, g, b, 1 ) );
+				computeInitNode = computeInit().compute( width * height );
 
-					}
-				`, [ rand2 ] );
+				// compute ping-pong: blur function using .load() for textureLoad
+				const blur = Fn( ( [ readTex, uv ] ) => {
 
-				computeInitNode = computeInitWGSL( { writeTex: storageTexture( pingTexture ), index: instanceIndex, seed } ).compute( width * height );
+					const c0 = readTex.load( uv.add( ivec2( - 1, 1 ) ) );
+					const c1 = readTex.load( uv.add( ivec2( - 1, - 1 ) ) );
+					const c2 = readTex.load( uv.add( ivec2( 0, 0 ) ) );
+					const c3 = readTex.load( uv.add( ivec2( 1, - 1 ) ) );
+					const c4 = readTex.load( uv.add( ivec2( 1, 1 ) ) );
 
-				// compute loop
+					return c0.add( c1 ).add( c2 ).add( c3 ).add( c4 ).div( 5.0 );
 
-				const computePingPongWGSL = wgslFn( `
-					fn computePingPongWGSL( readTex: texture_storage_2d<${wgslFormat}, read>, writeTex: texture_storage_2d<${ wgslFormat }, write>, index: u32 ) -> void {
+				} );
 
-						let posX = index % ${ width };
-						let posY = index / ${ width };
-						let indexUV = vec2i( i32( posX ), i32( posY ) );
+				// compute loop: read from one texture, blur, write to another
+				const computePingPong = Fn( ( [ readTex, writeTex ] ) => {
 
-						let color = blur( readTex, indexUV ).rgb;
+					const posX = instanceIndex.mod( width );
+					const posY = instanceIndex.div( width );
+					const indexUV = ivec2( int( posX ), int( posY ) );
 
-						textureStore( writeTex, indexUV, vec4f( color * 1.05, 1 ) );
+					const color = blur( readTex, indexUV );
 
-					}
-				`, [ rand2 ] );
+					textureStore( writeTex, indexUV, vec4( color.rgb.mul( 1.05 ), 1 ) );
 
-				//
+				} );
 
-				computeToPong = computePingPongWGSL( { readTex: readPing, writeTex: writePong, index: instanceIndex } ).compute( width * height );
-				computeToPing = computePingPongWGSL( { readTex: readPong, writeTex: writePing, index: instanceIndex } ).compute( width * height );
+				computeToPong = computePingPong( readPing, writePong ).compute( width * height );
+				computeToPing = computePingPong( readPong, writePing ).compute( width * height );
 
 				//
 

src/nodes/accessors/StorageTextureNode.js

@@ -222,6 +222,7 @@ class StorageTextureNode extends TextureNode {
 		const newNode = super.clone();
 		newNode.storeNode = this.storeNode;
 		newNode.mipLevel = this.mipLevel;
+		newNode.access = this.access;
 		return newNode;
 
 	}
@@ -255,7 +256,20 @@ export const storageTexture = /*@__PURE__*/ nodeProxy( StorageTextureNode ).setP
  */
 export const textureStore = ( value, uvNode, storeNode ) => {
 
-	const node = storageTexture( value, uvNode, storeNode );
+	let node;
+
+	if ( value.isStorageTextureNode === true ) {
+
+		// Derive new storage texture node from existing one
+		node = value.clone();
+		node.uvNode = uvNode;
+		node.storeNode = storeNode;
+
+	} else {
+
+		node = storageTexture( value, uvNode, storeNode );
+
+	}
 
 	if ( storeNode !== null ) node.toStack();
 

src/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -548,7 +548,9 @@ class WGSLNodeBuilder extends NodeBuilder {
 	 */
 	generateTextureLoad( texture, textureProperty, uvIndexSnippet, levelSnippet, depthSnippet, offsetSnippet ) {
 
-		if ( levelSnippet === null ) levelSnippet = '0u';
+		const isStorageTexture = texture.isStorageTexture === true;
+
+		if ( levelSnippet === null && ! isStorageTexture ) levelSnippet = '0u';
 
 		if ( offsetSnippet ) {
 
@@ -560,15 +562,33 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 		if ( depthSnippet ) {
 
-			snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, ${ depthSnippet }, u32( ${ levelSnippet } ) )`;
+			// Storage textures don't take a level parameter in WGSL
+			if ( isStorageTexture ) {
+
+				snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, ${ depthSnippet } )`;
+
+			} else {
+
+				snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, ${ depthSnippet }, u32( ${ levelSnippet } ) )`;
+
+			}
 
 		} else {
 
-			snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, u32( ${ levelSnippet } ) )`;
+			// Storage textures don't take a level parameter in WGSL
+			if ( isStorageTexture ) {
 
-			if ( this.renderer.backend.compatibilityMode && texture.isDepthTexture ) {
+				snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet } )`;
 
-				snippet += '.x';
+			} else {
+
+				snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, u32( ${ levelSnippet } ) )`;
+
+				if ( this.renderer.backend.compatibilityMode && texture.isDepthTexture ) {
+
+					snippet += '.x';
+
+				}
 
 			}
 

test/unit/src/nodes/accessors/StorageTextureNode.tests.js

@@ -0,0 +1,39 @@
+import { storageTexture } from '../../../../../src/nodes/accessors/StorageTextureNode.js';
+import { NodeAccess } from '../../../../../src/nodes/core/constants.js';
+import StorageTexture from '../../../../../src/renderers/common/StorageTexture.js';
+
+export default QUnit.module( 'Nodes', () => {
+
+	QUnit.module( 'Accessors', () => {
+
+		QUnit.module( 'StorageTextureNode', () => {
+
+			QUnit.test( 'clone preserves access property', ( assert ) => {
+
+				const texture = new StorageTexture( 512, 512 );
+				const node = storageTexture( texture ).setAccess( NodeAccess.READ_ONLY );
+
+				assert.strictEqual( node.access, NodeAccess.READ_ONLY, 'original has READ_ONLY access' );
+
+				const cloned = node.clone();
+
+				assert.strictEqual( cloned.access, NodeAccess.READ_ONLY, 'cloned node preserves READ_ONLY access' );
+
+			} );
+
+			QUnit.test( 'clone preserves READ_WRITE access', ( assert ) => {
+
+				const texture = new StorageTexture( 512, 512 );
+				const node = storageTexture( texture ).setAccess( NodeAccess.READ_WRITE );
+
+				const cloned = node.clone();
+
+				assert.strictEqual( cloned.access, NodeAccess.READ_WRITE, 'cloned node preserves READ_WRITE access' );
+
+			} );
+
+		} );
+
+	} );
+
+} );

test/unit/src/renderers/webgpu/nodes/WGSLNodeBuilder.tests.js

@@ -0,0 +1,69 @@
+import WGSLNodeBuilder from '../../../../../../src/renderers/webgpu/nodes/WGSLNodeBuilder.js';
+
+export default QUnit.module( 'Renderers', () => {
+
+	QUnit.module( 'WebGPU', () => {
+
+		QUnit.module( 'Nodes', () => {
+
+			QUnit.module( 'WGSLNodeBuilder', () => {
+
+				// generateTextureLoad is essentially a pure function (texture info -> WGSL string)
+				// The only 'this' access is renderer.backend.compatibilityMode for a depth texture edge case
+				// We test the real method with minimal context to verify WGSL output
+
+				QUnit.test( 'generateTextureLoad omits level for storage textures', ( assert ) => {
+
+					const context = {
+						renderer: { backend: { compatibilityMode: false } }
+					};
+
+					const storageTexture = { isStorageTexture: true };
+
+					const snippet = WGSLNodeBuilder.prototype.generateTextureLoad.call(
+						context,
+						storageTexture,
+						'testTexture',
+						'uvec2(0, 0)',
+						null, // levelSnippet
+						null, // depthSnippet
+						null // offsetSnippet
+					);
+
+					// Storage textures should NOT have level parameter (WGSL spec)
+					assert.notOk( snippet.includes( 'u32(' ), 'storage texture load should not include level parameter' );
+					assert.strictEqual( snippet, 'textureLoad( testTexture, uvec2(0, 0) )', 'correct WGSL for storage texture' );
+
+				} );
+
+				QUnit.test( 'generateTextureLoad includes level for regular textures', ( assert ) => {
+
+					const context = {
+						renderer: { backend: { compatibilityMode: false } }
+					};
+
+					const regularTexture = { isStorageTexture: false };
+
+					const snippet = WGSLNodeBuilder.prototype.generateTextureLoad.call(
+						context,
+						regularTexture,
+						'testTexture',
+						'uvec2(0, 0)',
+						null, // levelSnippet - should default to '0u'
+						null,
+						null
+					);
+
+					// Regular textures SHOULD have level parameter
+					assert.ok( snippet.includes( 'u32( 0u )' ), 'regular texture load should include default level parameter' );
+					assert.strictEqual( snippet, 'textureLoad( testTexture, uvec2(0, 0), u32( 0u ) )' );
+
+				} );
+
+			} );
+
+		} );
+
+	} );
+
+} );

test/unit/three.source.unit.js

@@ -265,6 +265,12 @@ import './src/renderers/webgl/WebGLTextures.tests.js';
 import './src/renderers/webgl/WebGLUniforms.tests.js';
 import './src/renderers/webgl/WebGLUtils.tests.js';
 
+//src/renderers/webgpu/nodes
+import './src/renderers/webgpu/nodes/WGSLNodeBuilder.tests.js';
+
+//src/nodes/accessors
+import './src/nodes/accessors/StorageTextureNode.tests.js';
+
 
 //src/scenes
 import './src/scenes/Fog.tests.js';