WebGPURenderer: `StorageBufferNode` Support reading external elements in the WebGL Backend (#27661)

* init pbo

* remove flag

* regenerate live example

* single buffer and alternate for example

* cleanup, no idea about circ dep

* test fix circular

* cleanup

* moved pbo management to GLSLNodeBuilder

* fix screenshots

* support more ranges

* format dynamically in arrayelementnode

* init vertex buffer allocation to prevent issues with some backends

* support different type size and update example

* fix files.json

* puppeteer need webgpu support

* unbind post pbo

* add TODO

* cleanup

* Move to StorageArrayElementNode

* cleanup

* rev

* add increaseUsage

* Update StorageArrayElementNode.js

* fixes optmization and revisions

* revision

* improved example for E2E tests

* WIP: Test (1)

* handle compute with non-PBO -> PBO case

* revert WIP

* TSL: add `storageObject`

* revision

* make sure attributeData gets pbo attached in GLSLNodeBuilder

* no need transfer anymore

* revert compute_texture

* remove copyBufferToSubBuffer

---------

Co-authored-by: sunag <[email protected]>

examples/files.json

@@ -374,7 +374,8 @@
 		"webgpu_postprocessing_anamorphic",
 		"webgpu_mirror",
 		"webgpu_multisampled_renderbuffers",
-		"webgpu_materials_texture_anisotropy"
+		"webgpu_materials_texture_anisotropy",
+		"webgpu_storage_buffer"
 	],
 	"webaudio": [
 		"webaudio_orientation",

examples/jsm/nodes/Nodes.js

@@ -98,7 +98,7 @@ export { default as ReferenceNode, reference, referenceIndex } from './accessors
 export { default as ReflectVectorNode, reflectVector } from './accessors/ReflectVectorNode.js';
 export { default as SkinningNode, skinning } from './accessors/SkinningNode.js';
 export { default as SceneNode, backgroundBlurriness, backgroundIntensity } from './accessors/SceneNode.js';
-export { default as StorageBufferNode, storage } from './accessors/StorageBufferNode.js';
+export { default as StorageBufferNode, storage, storageObject } from './accessors/StorageBufferNode.js';
 export { default as TangentNode, tangentGeometry, tangentLocal, tangentView, tangentWorld, transformedTangentView, transformedTangentWorld } from './accessors/TangentNode.js';
 export { default as TextureNode, texture, textureLoad, /*textureLevel,*/ sampler } from './accessors/TextureNode.js';
 export { default as TextureStoreNode, textureStore } from './accessors/TextureStoreNode.js';

examples/jsm/nodes/accessors/StorageBufferNode.js

@@ -13,6 +13,8 @@ class StorageBufferNode extends BufferNode {
 
 		this.isStorageBufferNode = true;
 
+		this.bufferObject = false;
+
 		this._attribute = null;
 		this._varying = null;
 
@@ -30,6 +32,14 @@ class StorageBufferNode extends BufferNode {
 
 	}
 
+	setBufferObject( value ) {
+
+		this.bufferObject = value;
+
+		return this;
+
+	}
+
 	generate( builder ) {
 
 		if ( builder.isAvailable( 'storageBuffer' ) ) return super.generate( builder );
@@ -57,5 +67,6 @@ class StorageBufferNode extends BufferNode {
 export default StorageBufferNode;
 
 export const storage = ( value, type, count ) => nodeObject( new StorageBufferNode( value, type, count ) );
+export const storageObject = ( value, type, count ) => nodeObject( new StorageBufferNode( value, type, count ).setBufferObject( true ) );
 
 addNodeClass( 'StorageBufferNode', StorageBufferNode );

examples/jsm/nodes/core/AssignNode.js

@@ -27,12 +27,11 @@ class AssignNode extends TempNode {
 
 	generate( builder, output ) {
 
-		const targetNode = this.targetNode;
-		const sourceNode = this.sourceNode;
+		const { targetNode, sourceNode } = this;
 
 		const targetType = targetNode.getNodeType( builder );
 
-		const target = targetNode.build( builder );
+		const target = targetNode.context( { assign: true } ).build( builder );
 		const source = sourceNode.build( builder, targetType );
 
 		const snippet = `${ target } = ${ source }`;

examples/jsm/nodes/utils/StorageArrayElementNode.js

@@ -24,13 +24,43 @@ class StorageArrayElementNode extends ArrayElementNode {
 
 	}
 
-	generate( builder ) {
+	setup( builder ) {
+
+		if ( builder.isAvailable( 'storageBuffer' ) === false ) {
+
+			if ( ! this.node.instanceIndex && this.node.bufferObject === true ) {
+
+				builder.setupPBO( this.node );
+
+			}
+
+		}
+
+		return super.setup( builder );
+
+	}
+
+	generate( builder, output ) {
 
 		let snippet;
 
+		const isAssignContext = builder.context.assign;
+
+		//
+
 		if ( builder.isAvailable( 'storageBuffer' ) === false ) {
 
-			snippet = this.node.build( builder );
+			const { node } = this;
+
+			if ( ! node.instanceIndex && this.node.bufferObject === true && isAssignContext !== true ) {
+
+				snippet = builder.generatePBO( this );
+
+			} else {
+
+				snippet = node.build( builder );
+
+			}
 
 		} else {
 
@@ -38,6 +68,14 @@ class StorageArrayElementNode extends ArrayElementNode {
 
 		}
 
+		if ( isAssignContext !== true ) {
+
+			const type = this.getNodeType( builder );
+
+			snippet = builder.format( snippet, type, output );
+
+		}
+
 		return snippet;
 
 	}

examples/jsm/renderers/webgl/WebGLBackend.js

@@ -428,9 +428,20 @@ class WebGLBackend extends Backend {
 		gl.bindTransformFeedback( gl.TRANSFORM_FEEDBACK, null );
 
 		// switch active buffers
+
 		for ( let i = 0; i < transformBuffers.length; i ++ ) {
 
-			transformBuffers[ i ].switchBuffers();
+			const dualAttributeData = transformBuffers[ i ];
+
+			if ( dualAttributeData.pbo ) {
+
+				this.textureUtils.copyBufferToTexture( dualAttributeData.transformBuffer, dualAttributeData.pbo );
+
+			} else {
+
+				dualAttributeData.switchBuffers();
+
+			}
 
 		}
 

examples/jsm/renderers/webgl/nodes/GLSLNodeBuilder.js

@@ -1,11 +1,12 @@
-import { MathNode, GLSLNodeParser, NodeBuilder } from '../../../nodes/Nodes.js';
+import { MathNode, GLSLNodeParser, NodeBuilder, UniformNode, vectorComponents } from '../../../nodes/Nodes.js';
 
 import UniformBuffer from '../../common/UniformBuffer.js';
 import NodeUniformsGroup from '../../common/nodes/NodeUniformsGroup.js';
 
 import { NodeSampledTexture, NodeSampledCubeTexture } from '../../common/nodes/NodeSampledTexture.js';
 
-import { IntType } from 'three';
+
+import { IntType, DataTexture, RGBAFormat, FloatType } from 'three';
 
 const glslMethods = {
 	[ MathNode.ATAN2 ]: 'atan',
@@ -84,6 +85,131 @@ ${ flowData.code }
 
 	}
 
+	setupPBO( storageBufferNode ) {
+
+		const attribute = storageBufferNode.value;
+
+		if ( attribute.pbo === undefined ) {
+
+			const originalArray = attribute.array;
+			const numElements = attribute.count * attribute.itemSize;
+
+			const width = Math.pow( 2, Math.ceil( Math.log2( Math.sqrt( numElements / 4 ) ) ) );
+			let height = Math.ceil( ( numElements / 4 ) / width );
+			if ( width * height * 4 < numElements ) height ++; // Ensure enough space
+
+			const newSize = width * height * 4; // 4 floats per pixel due to RGBA format
+
+			const newArray = new Float32Array( newSize );
+
+			newArray.set( originalArray, 0 );
+
+			attribute.array = newArray;
+			attribute.count = newSize;
+
+			const pboTexture = new DataTexture( attribute.array, width, height, RGBAFormat, FloatType );
+			pboTexture.needsUpdate = true;
+			pboTexture.isPBOTexture = true;
+
+			const pbo = new UniformNode( pboTexture );
+			pbo.setPrecision( 'high' );
+
+			attribute.pboNode = pbo;
+			attribute.pbo = pbo.value;
+
+			this.getUniformFromNode( attribute.pboNode, 'texture', this.shaderStage, this.context.label );
+
+		}
+
+	}
+
+	generatePBO( storageArrayElementNode ) {
+
+		const { node, indexNode } = storageArrayElementNode;
+		const attribute = node.value;
+
+		if ( this.renderer.backend.has( attribute ) ) {
+
+			const attributeData = this.renderer.backend.get( attribute );
+			attributeData.pbo = attribute.pbo;
+
+		}
+
+
+		const nodeUniform = this.getUniformFromNode( attribute.pboNode, 'texture', this.shaderStage, this.context.label );
+		const textureName = this.getPropertyName( nodeUniform );
+
+		indexNode.increaseUsage( this ); // force cache generate to be used as index in x,y
+		const indexSnippet = indexNode.build( this, 'uint' );
+
+		const elementNodeData = this.getDataFromNode( storageArrayElementNode );
+
+		let propertyName = elementNodeData.propertyName;
+
+		if ( propertyName === undefined ) {
+
+			// property element
+
+			const nodeVar = this.getVarFromNode( storageArrayElementNode );
+
+			propertyName = this.getPropertyName( nodeVar );
+
+			// property size
+
+			const bufferNodeData = this.getDataFromNode( node );
+
+			let propertySizeName = bufferNodeData.propertySizeName;
+
+			if ( propertySizeName === undefined ) {
+
+				propertySizeName = propertyName + 'Size';
+
+				this.getVarFromNode( node, propertySizeName, 'uint' );
+
+				this.addLineFlowCode( `${ propertySizeName } = uint( textureSize( ${ textureName }, 0 ).x )` );
+
+				bufferNodeData.propertySizeName = propertySizeName;
+
+			}
+
+			//
+
+			let channel;
+			let padding;
+
+			const itemSize = attribute.itemSize;
+
+			if ( itemSize === 1 ) {
+
+				padding = 4;
+				channel = `[ ${indexSnippet} % uint( ${ padding } ) ]`;
+
+			} else {
+
+				padding = itemSize > 2 ? 1 : itemSize;
+				channel = '.' + vectorComponents.join( '' ).slice( 0, itemSize );
+
+			}
+
+			const uvSnippet = `ivec2(
+				${indexSnippet} / uint( ${ padding } ) % ${ propertySizeName },
+				${indexSnippet} / ( uint( ${ padding } ) * ${ propertySizeName } )
+			)`;
+
+			const snippet = this.generateTextureLoad( null, textureName, uvSnippet, null, '0' );
+
+			//
+
+			this.addLineFlowCode( `${ propertyName } = ${ snippet + channel }` );
+
+			elementNodeData.propertyName = propertyName;
+
+		}
+
+		return propertyName;
+
+	}
+
 	generateTextureLoad( texture, textureProperty, uvIndexSnippet, depthSnippet, levelSnippet = '0' ) {
 
 		if ( depthSnippet ) {

examples/jsm/renderers/webgl/utils/WebGLAttributeUtils.js

@@ -8,6 +8,8 @@ class DualAttributeData {
 
 		this.buffers = [ attributeData.bufferGPU, dualBuffer ];
 		this.type = attributeData.type;
+		this.pbo = attributeData.pbo;
+		this.byteLength = attributeData.byteLength;
 		this.bytesPerElement = attributeData.BYTES_PER_ELEMENT;
 		this.version = attributeData.version;
 		this.isInteger = attributeData.isInteger;
@@ -127,8 +129,10 @@ class WebGLAttributeUtils {
 		let attributeData = {
 			bufferGPU,
 			type,
+			byteLength: array.byteLength,
 			bytesPerElement: array.BYTES_PER_ELEMENT,
 			version: attribute.version,
+			pbo: attribute.pbo,
 			isInteger: type === gl.INT || type === gl.UNSIGNED_INT || attribute.gpuType === IntType,
 			id: _id ++
 		};

examples/jsm/renderers/webgl/utils/WebGLTextureUtils.js

@@ -207,6 +207,7 @@ class WebGLTextureUtils {
 			if ( texture.type === FloatType && extensions.has( 'OES_texture_float_linear' ) === false ) return; // verify extension for WebGL 1 and WebGL 2
 
 			if ( texture.anisotropy > 1 || currentAnisotropy !== texture.anisotropy ) {
+
 				const extension = extensions.get( 'EXT_texture_filter_anisotropic' );
 				gl.texParameterf( textureType, extension.TEXTURE_MAX_ANISOTROPY_EXT, Math.min( texture.anisotropy, backend.getMaxAnisotropy() ) );
 				backend.get( texture ).currentAnisotropy = texture.anisotropy;
@@ -289,6 +290,41 @@ class WebGLTextureUtils {
 
 	}
 
+	copyBufferToTexture( buffer, texture ) {
+
+		const { gl, backend } = this;
+
+		const { textureGPU, glTextureType, glFormat, glType } = backend.get( texture );
+
+		const { width, height } = texture.source.data;
+
+		gl.bindBuffer( gl.PIXEL_UNPACK_BUFFER, buffer );
+
+		backend.state.bindTexture( glTextureType, textureGPU );
+
+		gl.pixelStorei( gl.UNPACK_FLIP_Y_WEBGL, false );
+		gl.pixelStorei( gl.UNPACK_PREMULTIPLY_ALPHA_WEBGL, false );
+		gl.texSubImage2D( glTextureType, 0, 0, 0, width, height, glFormat, glType, 0 );
+
+		gl.bindBuffer( gl.PIXEL_UNPACK_BUFFER, null );
+
+		backend.state.unbindTexture();
+		// debug
+		// const framebuffer = gl.createFramebuffer();
+		// gl.bindFramebuffer( gl.FRAMEBUFFER, framebuffer );
+		// gl.framebufferTexture2D( gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, glTextureType, textureGPU, 0 );
+
+		// const readout = new Float32Array( width * height * 4 );
+
+		// const altFormat = gl.getParameter( gl.IMPLEMENTATION_COLOR_READ_FORMAT );
+		// const altType = gl.getParameter( gl.IMPLEMENTATION_COLOR_READ_TYPE );
+
+		// gl.readPixels( 0, 0, width, height, altFormat, altType, readout );
+		// gl.bindFramebuffer( gl.FRAMEBUFFER, null );
+		// console.log( readout );
+
+	}
+
 	updateTexture( texture, options ) {
 
 		const { gl } = this;

examples/webgpu_compute_audio.html

@@ -29,11 +29,10 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { tslFn, uniform, storage, instanceIndex, float, texture, viewportTopLeft, color } from 'three/nodes';
+			import { tslFn, uniform, storage, storageObject, instanceIndex, float, texture, viewportTopLeft, color } from 'three/nodes';
 
 			import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
 
-			import WebGPU from 'three/addons/capabilities/WebGPU.js';
 			import WebGPURenderer from 'three/addons/renderers/webgpu/WebGPURenderer.js';
 			import StorageBufferAttribute from 'three/addons/renderers/common/StorageBufferAttribute.js';
 
@@ -53,7 +52,7 @@
 
 				// compute audio
 
-				renderer.compute( computeNode );
+				await renderer.computeAsync( computeNode );
 
 				const waveArray = new Float32Array( await renderer.getArrayBufferAsync( waveGPUBuffer ) );
 
@@ -82,22 +81,15 @@
 
 			async function init() {
 
-				if ( WebGPU.isAvailable() === false ) {
+				// if ( WebGPU.isAvailable() === false ) {
 
-					document.body.appendChild( WebGPU.getErrorMessage() );
+				// 	document.body.appendChild( WebGPU.getErrorMessage() );
 
-					throw new Error( 'No WebGPU support' );
+				// 	throw new Error( 'No WebGPU support' );
 
-				}
-
-				document.onclick = () => {
-
-					const overlay = document.getElementById( 'overlay' );
-					if ( overlay !== null ) overlay.remove();
+				// }
 
-					playAudioBuffer();
-
-				};
+			
 
 				// audio buffer
 
@@ -123,7 +115,7 @@
 
 				// read-only buffer
 
-				const waveNode = storage( new StorageBufferAttribute( waveBuffer, 1 ), 'float', waveBuffer.length );
+				const waveNode = storageObject( new StorageBufferAttribute( waveBuffer, 1 ), 'float', waveBuffer.length );
 
 
 				// params
@@ -212,6 +204,16 @@
 
 				window.addEventListener( 'resize', onWindowResize );
 
+
+				document.onclick = () => {
+
+					const overlay = document.getElementById( 'overlay' );
+					if ( overlay !== null ) overlay.remove();
+
+					playAudioBuffer();
+
+				};
+
 			}
 
 			function onWindowResize() {

examples/webgpu_storage_buffer.html

@@ -0,0 +1,158 @@
+<html lang="en">
+	<head>
+		<title>three.js - WebGPU - Storage PBO External Element</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>
+			<br />This example demonstrate the fetch of external element from a StorageBuffer.
+			<br /> Left canvas is using WebGPU Backend, right canvas is WebGL Backend.
+		</div>
+
+		<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 { storageObject, vec3, uv, uint, float, tslFn, instanceIndex, MeshBasicNodeMaterial } from 'three/nodes';
+
+			import WebGPURenderer from 'three/addons/renderers/webgpu/WebGPURenderer.js';
+			import StorageBufferAttribute from 'three/addons/renderers/common/StorageBufferAttribute.js';
+
+			// WebGPU Backend
+			init();
+
+			// WebGL Backend
+			init( true );
+
+			function init( forceWebGL = false ) {
+
+				const aspect = ( window.innerWidth / 2 ) / window.innerHeight;
+				const camera = new THREE.OrthographicCamera( - aspect, aspect, 1, - 1, 0, 2 );
+				camera.position.z = 1;
+
+				const scene = new THREE.Scene();
+
+				// texture
+
+				const typeSize = 1; // 1:'float', 2:'vec2', 4:'vec4' -> use power of 2
+				const size = 1024;
+
+				const array = new Array( size * typeSize ).fill( 0 );
+
+				const type = [ 'float', 'vec2', 'vec3', 'vec4' ][ typeSize - 1 ];
+
+				const arrayBuffer = new StorageBufferAttribute( new Float32Array( array ), typeSize );
+
+				const arrayBufferNode = storageObject( arrayBuffer, type, size );
+
+
+				const computeInitOrder = tslFn( () => {
+
+					arrayBufferNode.element( instanceIndex ).assign( uint( instanceIndex.div( typeSize ) ) );
+
+				} );
+
+				const computeInvertOrder = tslFn( () => {
+
+					const invertIndex = arrayBufferNode.element( float( size ).sub( instanceIndex ) );
+					arrayBufferNode.element( instanceIndex ).assign( invertIndex );
+
+				} );
+
+				// compute
+
+				const computeInit = computeInitOrder().compute( size );
+
+				const compute = computeInvertOrder().compute( size );
+
+				const material = new MeshBasicNodeMaterial( { color: 0x00ff00 } );
+
+				material.colorNode = tslFn( () => {
+
+					const index = uint( uv().x.mul( float( size ) ) );
+					const indexValue = arrayBufferNode.element( index ).toVar();
+					const value = float( indexValue ).div( float( size ) ).mul( 20 ).floor().div( 20 );
+
+					return vec3( value, value, value );
+
+				} )();
+			
+				//
+
+				const plane = new THREE.Mesh( new THREE.PlaneGeometry( 1, 1 ), material );
+				scene.add( plane );
+
+				const renderer = new WebGPURenderer( { antialias: false, forceWebGL: forceWebGL } );
+				renderer.setPixelRatio( window.devicePixelRatio );
+				renderer.setSize( window.innerWidth / 2, window.innerHeight );
+
+				document.body.appendChild( renderer.domElement );
+				renderer.domElement.style.position = 'absolute';
+				renderer.domElement.style.top = '0';
+				renderer.domElement.style.left = '0';
+				renderer.domElement.style.width = '50%';
+				renderer.domElement.style.height = '100%';
+
+				if ( forceWebGL ) {
+
+					renderer.domElement.style.left = '50%';
+
+					scene.background = new THREE.Color( 0x212121 );
+			
+				} else {
+
+					scene.background = new THREE.Color( 0x313131 );
+
+				}
+
+				// Init Positions
+				renderer.compute( computeInit );
+
+				 const stepAnimation = async function () {
+
+					await renderer.computeAsync( compute );
+					await renderer.renderAsync( scene, camera );
+
+					setTimeout( stepAnimation, 1000 );
+			
+				};
+
+				stepAnimation();
+
+				window.addEventListener( 'resize', onWindowResize );
+
+				function onWindowResize() {
+
+					renderer.setSize( window.innerWidth / 2, window.innerHeight );
+
+					const aspect = ( window.innerWidth / 2 ) / window.innerHeight;
+
+					const frustumHeight = camera.top - camera.bottom;
+
+					camera.left = - frustumHeight * aspect / 2;
+					camera.right = frustumHeight * aspect / 2;
+
+					camera.updateProjectionMatrix();
+
+					renderer.render( scene, camera );
+
+				}
+
+			}
+
+		</script>
+	</body>
+</html>

test/e2e/puppeteer.js

@@ -119,6 +119,9 @@ const exceptionList = [
 	'webgpu_sprites',
 	'webgpu_video_panorama',
 
+	// Awaiting for WebGPU Backend support in Puppeteer
+	'webgpu_storage_buffer',
+
 	// WebGPURenderer: Unknown problem
 	'webgpu_postprocessing_afterimage',
 	'webgpu_backdrop_water',