Skip to content

Commit 83b7ea5

Browse files
NateSmythNateSmyth
authored
StorageTextureNode: Add TSL read/write support (mrdoob#32734)
1 parent 3ce02cc commit 83b7ea5

6 files changed

Lines changed: 196 additions & 67 deletions

File tree

‎examples/webgpu_compute_texture_pingpong.html‎

Lines changed: 42 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
</div>
1717

1818
<small>
19-
Compute ping/pong texture using GPU.
19+
Compute ping/pong texture using GPU (pure TSL).
2020
</small>
2121
</div>
2222

@@ -34,7 +34,7 @@
3434
<script type="module">
3535

3636
import * as THREE from 'three/webgpu';
37-
import { storageTexture, wgslFn, code, instanceIndex, uniform, NodeAccess } from 'three/tsl';
37+
import { storageTexture, textureStore, Fn, instanceIndex, uniform, float, vec2, vec4, uvec2, ivec2, int, NodeAccess } from 'three/tsl';
3838

3939
import WebGPU from 'three/addons/capabilities/WebGPU.js';
4040

@@ -45,6 +45,8 @@
4545
let phase = true;
4646
let lastUpdate = - 1;
4747

48+
const width = 512, height = 512;
49+
4850
const seed = uniform( new THREE.Vector2() );
4951

5052
init();
@@ -68,7 +70,6 @@
6870
// texture
6971

7072
const hdr = true;
71-
const width = 512, height = 512;
7273

7374
pingTexture = new THREE.StorageTexture( width, height );
7475
pongTexture = new THREE.StorageTexture( width, height );
@@ -80,83 +81,63 @@
8081

8182
}
8283

83-
const wgslFormat = hdr ? 'rgba16float' : 'rgba8unorm';
84-
85-
const readPing = storageTexture( pingTexture ).setAccess( NodeAccess.READ_ONLY );
86-
const writePing = storageTexture( pingTexture ).setAccess( NodeAccess.WRITE_ONLY );
87-
const readPong = storageTexture( pongTexture ).setAccess( NodeAccess.READ_ONLY );
88-
const writePong = storageTexture( pongTexture ).setAccess( NodeAccess.WRITE_ONLY );
89-
90-
// compute init
91-
92-
const rand2 = code( `
93-
fn rand2( n: vec2f ) -> f32 {
94-
95-
return fract( sin( dot( n, vec2f( 12.9898, 4.1414 ) ) ) * 43758.5453 );
84+
const rand2 = Fn( ( [ n ] ) => {
9685

97-
}
86+
return n.dot( vec2( 12.9898, 4.1414 ) ).sin().mul( 43758.5453 ).fract();
9887

99-
fn blur( image : texture_storage_2d<${wgslFormat}, read>, uv : vec2i ) -> vec4f {
88+
} );
10089

101-
var color = vec4f( 0.0 );
102-
103-
color += textureLoad( image, uv + vec2i( - 1, 1 ));
104-
color += textureLoad( image, uv + vec2i( - 1, - 1 ));
105-
color += textureLoad( image, uv + vec2i( 0, 0 ));
106-
color += textureLoad( image, uv + vec2i( 1, - 1 ));
107-
color += textureLoad( image, uv + vec2i( 1, 1 ));
108-
109-
return color / 5.0;
110-
}
111-
112-
fn getUV( posX: u32, posY: u32 ) -> vec2f {
113-
114-
let uv = vec2f( f32( posX ) / ${ width }.0, f32( posY ) / ${ height }.0 );
90+
// Create storage texture nodes with proper access
91+
const writePing = storageTexture( pingTexture ).setAccess( NodeAccess.WRITE_ONLY );
92+
const readPing = storageTexture( pingTexture ).setAccess( NodeAccess.READ_ONLY );
93+
const writePong = storageTexture( pongTexture ).setAccess( NodeAccess.WRITE_ONLY );
94+
const readPong = storageTexture( pongTexture ).setAccess( NodeAccess.READ_ONLY );
11595

116-
return uv;
96+
const computeInit = Fn( () => {
11797

118-
}
119-
` );
98+
const posX = instanceIndex.mod( width );
99+
const posY = instanceIndex.div( width );
100+
const indexUV = uvec2( posX, posY );
101+
const uv = vec2( float( posX ).div( width ), float( posY ).div( height ) );
120102

121-
const computeInitWGSL = wgslFn( `
122-
fn computeInitWGSL( writeTex: texture_storage_2d<${ wgslFormat }, write>, index: u32, seed: vec2f ) -> void {
103+
const r = rand2( uv.add( seed.mul( 100 ) ) ).sub( rand2( uv.add( seed.mul( 300 ) ) ) );
104+
const g = rand2( uv.add( seed.mul( 200 ) ) ).sub( rand2( uv.add( seed.mul( 300 ) ) ) );
105+
const b = rand2( uv.add( seed.mul( 200 ) ) ).sub( rand2( uv.add( seed.mul( 100 ) ) ) );
123106

124-
let posX = index % ${ width };
125-
let posY = index / ${ width };
126-
let indexUV = vec2u( posX, posY );
127-
let uv = getUV( posX, posY );
107+
textureStore( writePing, indexUV, vec4( r, g, b, 1 ) );
128108

129-
let r = rand2( uv + seed * 100 ) - rand2( uv + seed * 300 );
130-
let g = rand2( uv + seed * 200 ) - rand2( uv + seed * 300 );
131-
let b = rand2( uv + seed * 200 ) - rand2( uv + seed * 100 );
109+
} );
132110

133-
textureStore( writeTex, indexUV, vec4( r, g, b, 1 ) );
111+
computeInitNode = computeInit().compute( width * height );
134112

135-
}
136-
`, [ rand2 ] );
113+
// compute ping-pong: blur function using .load() for textureLoad
114+
const blur = Fn( ( [ readTex, uv ] ) => {
137115

138-
computeInitNode = computeInitWGSL( { writeTex: storageTexture( pingTexture ), index: instanceIndex, seed } ).compute( width * height );
116+
const c0 = readTex.load( uv.add( ivec2( - 1, 1 ) ) );
117+
const c1 = readTex.load( uv.add( ivec2( - 1, - 1 ) ) );
118+
const c2 = readTex.load( uv.add( ivec2( 0, 0 ) ) );
119+
const c3 = readTex.load( uv.add( ivec2( 1, - 1 ) ) );
120+
const c4 = readTex.load( uv.add( ivec2( 1, 1 ) ) );
139121

140-
// compute loop
122+
return c0.add( c1 ).add( c2 ).add( c3 ).add( c4 ).div( 5.0 );
141123

142-
const computePingPongWGSL = wgslFn( `
143-
fn computePingPongWGSL( readTex: texture_storage_2d<${wgslFormat}, read>, writeTex: texture_storage_2d<${ wgslFormat }, write>, index: u32 ) -> void {
124+
} );
144125

145-
let posX = index % ${ width };
146-
let posY = index / ${ width };
147-
let indexUV = vec2i( i32( posX ), i32( posY ) );
126+
// compute loop: read from one texture, blur, write to another
127+
const computePingPong = Fn( ( [ readTex, writeTex ] ) => {
148128

149-
let color = blur( readTex, indexUV ).rgb;
129+
const posX = instanceIndex.mod( width );
130+
const posY = instanceIndex.div( width );
131+
const indexUV = ivec2( int( posX ), int( posY ) );
150132

151-
textureStore( writeTex, indexUV, vec4f( color * 1.05, 1 ) );
133+
const color = blur( readTex, indexUV );
152134

153-
}
154-
`, [ rand2 ] );
135+
textureStore( writeTex, indexUV, vec4( color.rgb.mul( 1.05 ), 1 ) );
155136

156-
//
137+
} );
157138

158-
computeToPong = computePingPongWGSL( { readTex: readPing, writeTex: writePong, index: instanceIndex } ).compute( width * height );
159-
computeToPing = computePingPongWGSL( { readTex: readPong, writeTex: writePing, index: instanceIndex } ).compute( width * height );
139+
computeToPong = computePingPong( readPing, writePong ).compute( width * height );
140+
computeToPing = computePingPong( readPong, writePing ).compute( width * height );
160141

161142
//
162143

‎src/nodes/accessors/StorageTextureNode.js‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,7 @@ class StorageTextureNode extends TextureNode {
222222
const newNode = super.clone();
223223
newNode.storeNode = this.storeNode;
224224
newNode.mipLevel = this.mipLevel;
225+
newNode.access = this.access;
225226
return newNode;
226227

227228
}
@@ -255,7 +256,20 @@ export const storageTexture = /*@__PURE__*/ nodeProxy( StorageTextureNode ).setP
255256
*/
256257
export const textureStore = ( value, uvNode, storeNode ) => {
257258

258-
const node = storageTexture( value, uvNode, storeNode );
259+
let node;
260+
261+
if ( value.isStorageTextureNode === true ) {
262+
263+
// Derive new storage texture node from existing one
264+
node = value.clone();
265+
node.uvNode = uvNode;
266+
node.storeNode = storeNode;
267+
268+
} else {
269+
270+
node = storageTexture( value, uvNode, storeNode );
271+
272+
}
259273

260274
if ( storeNode !== null ) node.toStack();
261275

‎src/renderers/webgpu/nodes/WGSLNodeBuilder.js‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -548,7 +548,9 @@ class WGSLNodeBuilder extends NodeBuilder {
548548
*/
549549
generateTextureLoad( texture, textureProperty, uvIndexSnippet, levelSnippet, depthSnippet, offsetSnippet ) {
550550

551-
if ( levelSnippet === null ) levelSnippet = '0u';
551+
const isStorageTexture = texture.isStorageTexture === true;
552+
553+
if ( levelSnippet === null && ! isStorageTexture ) levelSnippet = '0u';
552554

553555
if ( offsetSnippet ) {
554556

@@ -560,15 +562,33 @@ class WGSLNodeBuilder extends NodeBuilder {
560562

561563
if ( depthSnippet ) {
562564

563-
snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, ${ depthSnippet }, u32( ${ levelSnippet } ) )`;
565+
// Storage textures don't take a level parameter in WGSL
566+
if ( isStorageTexture ) {
567+
568+
snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, ${ depthSnippet } )`;
569+
570+
} else {
571+
572+
snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, ${ depthSnippet }, u32( ${ levelSnippet } ) )`;
573+
574+
}
564575

565576
} else {
566577

567-
snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, u32( ${ levelSnippet } ) )`;
578+
// Storage textures don't take a level parameter in WGSL
579+
if ( isStorageTexture ) {
568580

569-
if ( this.renderer.backend.compatibilityMode && texture.isDepthTexture ) {
581+
snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet } )`;
570582

571-
snippet += '.x';
583+
} else {
584+
585+
snippet = `textureLoad( ${ textureProperty }, ${ uvIndexSnippet }, u32( ${ levelSnippet } ) )`;
586+
587+
if ( this.renderer.backend.compatibilityMode && texture.isDepthTexture ) {
588+
589+
snippet += '.x';
590+
591+
}
572592

573593
}
574594

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import { storageTexture } from '../../../../../src/nodes/accessors/StorageTextureNode.js';
2+
import { NodeAccess } from '../../../../../src/nodes/core/constants.js';
3+
import StorageTexture from '../../../../../src/renderers/common/StorageTexture.js';
4+
5+
export default QUnit.module( 'Nodes', () => {
6+
7+
QUnit.module( 'Accessors', () => {
8+
9+
QUnit.module( 'StorageTextureNode', () => {
10+
11+
QUnit.test( 'clone preserves access property', ( assert ) => {
12+
13+
const texture = new StorageTexture( 512, 512 );
14+
const node = storageTexture( texture ).setAccess( NodeAccess.READ_ONLY );
15+
16+
assert.strictEqual( node.access, NodeAccess.READ_ONLY, 'original has READ_ONLY access' );
17+
18+
const cloned = node.clone();
19+
20+
assert.strictEqual( cloned.access, NodeAccess.READ_ONLY, 'cloned node preserves READ_ONLY access' );
21+
22+
} );
23+
24+
QUnit.test( 'clone preserves READ_WRITE access', ( assert ) => {
25+
26+
const texture = new StorageTexture( 512, 512 );
27+
const node = storageTexture( texture ).setAccess( NodeAccess.READ_WRITE );
28+
29+
const cloned = node.clone();
30+
31+
assert.strictEqual( cloned.access, NodeAccess.READ_WRITE, 'cloned node preserves READ_WRITE access' );
32+
33+
} );
34+
35+
} );
36+
37+
} );
38+
39+
} );

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
import WGSLNodeBuilder from '../../../../../../src/renderers/webgpu/nodes/WGSLNodeBuilder.js';
2+
3+
export default QUnit.module( 'Renderers', () => {
4+
5+
QUnit.module( 'WebGPU', () => {
6+
7+
QUnit.module( 'Nodes', () => {
8+
9+
QUnit.module( 'WGSLNodeBuilder', () => {
10+
11+
// generateTextureLoad is essentially a pure function (texture info -> WGSL string)
12+
// The only 'this' access is renderer.backend.compatibilityMode for a depth texture edge case
13+
// We test the real method with minimal context to verify WGSL output
14+
15+
QUnit.test( 'generateTextureLoad omits level for storage textures', ( assert ) => {
16+
17+
const context = {
18+
renderer: { backend: { compatibilityMode: false } }
19+
};
20+
21+
const storageTexture = { isStorageTexture: true };
22+
23+
const snippet = WGSLNodeBuilder.prototype.generateTextureLoad.call(
24+
context,
25+
storageTexture,
26+
'testTexture',
27+
'uvec2(0, 0)',
28+
null, // levelSnippet
29+
null, // depthSnippet
30+
null // offsetSnippet
31+
);
32+
33+
// Storage textures should NOT have level parameter (WGSL spec)
34+
assert.notOk( snippet.includes( 'u32(' ), 'storage texture load should not include level parameter' );
35+
assert.strictEqual( snippet, 'textureLoad( testTexture, uvec2(0, 0) )', 'correct WGSL for storage texture' );
36+
37+
} );
38+
39+
QUnit.test( 'generateTextureLoad includes level for regular textures', ( assert ) => {
40+
41+
const context = {
42+
renderer: { backend: { compatibilityMode: false } }
43+
};
44+
45+
const regularTexture = { isStorageTexture: false };
46+
47+
const snippet = WGSLNodeBuilder.prototype.generateTextureLoad.call(
48+
context,
49+
regularTexture,
50+
'testTexture',
51+
'uvec2(0, 0)',
52+
null, // levelSnippet - should default to '0u'
53+
null,
54+
null
55+
);
56+
57+
// Regular textures SHOULD have level parameter
58+
assert.ok( snippet.includes( 'u32( 0u )' ), 'regular texture load should include default level parameter' );
59+
assert.strictEqual( snippet, 'textureLoad( testTexture, uvec2(0, 0), u32( 0u ) )' );
60+
61+
} );
62+
63+
} );
64+
65+
} );
66+
67+
} );
68+
69+
} );

‎test/unit/three.source.unit.js‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,12 @@ import './src/renderers/webgl/WebGLTextures.tests.js';
265265
import './src/renderers/webgl/WebGLUniforms.tests.js';
266266
import './src/renderers/webgl/WebGLUtils.tests.js';
267267

268+
//src/renderers/webgpu/nodes
269+
import './src/renderers/webgpu/nodes/WGSLNodeBuilder.tests.js';
270+
271+
//src/nodes/accessors
272+
import './src/nodes/accessors/StorageTextureNode.tests.js';
273+
268274

269275
//src/scenes
270276
import './src/scenes/Fog.tests.js';

0 commit comments

Comments
 (0)