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 <sunagbrasil@gmail.com>

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,20 +24,58 @@ 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 {
 
 			snippet = super.generate( builder );
 
 		}
 
+		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;