WebGPURenderer: Fix `vec2` and `vec3` for `storageObject` in StorageBufferNode (#27697)

* improve tests

* fix vec2, vec3 packing

* cleanup

* always switch buffer to support reading and align with toAttribute()

examples/jsm/renderers/webgl/WebGLBackend.js

@@ -437,11 +437,10 @@ class WebGLBackend extends Backend {
 
 				this.textureUtils.copyBufferToTexture( dualAttributeData.transformBuffer, dualAttributeData.pbo );
 
-			} else {
+			}
 
-				dualAttributeData.switchBuffers();
+			dualAttributeData.switchBuffers();
 
-			}
 
 		}
 

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

@@ -6,7 +6,7 @@ import NodeUniformsGroup from '../../common/nodes/NodeUniformsGroup.js';
 import { NodeSampledTexture, NodeSampledCubeTexture } from '../../common/nodes/NodeSampledTexture.js';
 
 
-import { IntType, DataTexture, RGBAFormat, FloatType } from 'three';
+import { RedFormat, RGFormat, IntType, DataTexture, RGBAFormat, FloatType } from 'three';
 
 const glslMethods = {
 	[ MathNode.ATAN2 ]: 'atan',
@@ -94,20 +94,36 @@ ${ flowData.code }
 			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 { itemSize } = attribute;
+			let format = RedFormat;
 
-			const newSize = width * height * 4; // 4 floats per pixel due to RGBA format
+			if ( itemSize === 2 ) {
+
+				format = RGFormat;
+
+			} else if ( itemSize === 3 ) {
+
+				format = 6407; // patch since legacy doesn't use RGBFormat for rendering but here it's needed for packing optimization
+
+			} else if ( itemSize === 4 ) {
+
+				format = RGBAFormat;
+
+			}
+
+			const width = Math.pow( 2, Math.ceil( Math.log2( Math.sqrt( numElements / itemSize ) ) ) );
+			let height = Math.ceil( ( numElements / itemSize ) / width );
+			if ( width * height * itemSize < numElements ) height ++; // Ensure enough space
+
+			const newSize = width * height * itemSize;
 
 			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 );
+			const pboTexture = new DataTexture( attribute.array, width, height, format, FloatType );
 			pboTexture.needsUpdate = true;
 			pboTexture.isPBOTexture = true;
 
@@ -174,27 +190,10 @@ ${ flowData.code }
 
 			//
 
-			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 { itemSize } = attribute;
 
-			const uvSnippet = `ivec2(
-				${indexSnippet} / uint( ${ padding } ) % ${ propertySizeName },
-				${indexSnippet} / ( uint( ${ padding } ) * ${ propertySizeName } )
-			)`;
+			const channel = '.' + vectorComponents.join( '' ).slice( 0, itemSize );
+			const uvSnippet = `ivec2(${indexSnippet} % ${ propertySizeName }, ${indexSnippet} / ${ propertySizeName })`;
 
 			const snippet = this.generateTextureLoad( null, textureName, uvSnippet, null, '0' );
 

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

@@ -135,6 +135,17 @@ class WebGLTextureUtils {
 
 		}
 
+		if ( glFormat === gl.RGB ) {
+
+			if ( glType === gl.FLOAT ) internalFormat = gl.RGB32F;
+			if ( glType === gl.HALF_FLOAT ) internalFormat = gl.RGB16F;
+			if ( glType === gl.UNSIGNED_BYTE ) internalFormat = gl.RGB8;
+			if ( glType === gl.UNSIGNED_SHORT_5_6_5 ) internalFormat = gl.RGB565;
+			if ( glType === gl.UNSIGNED_SHORT_5_5_5_1 ) internalFormat = gl.RGB5_A1;
+			if ( glType === gl.UNSIGNED_SHORT_4_4_4_4 ) internalFormat = gl.RGB4;
+
+		}
+
 		if ( glFormat === gl.RGBA ) {
 
 			if ( glType === gl.FLOAT ) internalFormat = gl.RGBA32F;

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

@@ -35,6 +35,7 @@ class WebGLUtils {
 		}
 
 		if ( p === AlphaFormat ) return gl.ALPHA;
+		if ( p === gl.RGB ) return gl.RGB; // patch since legacy doesn't use RGBFormat for rendering but here it's needed for packing optimization
 		if ( p === RGBAFormat ) return gl.RGBA;
 		if ( p === LuminanceFormat ) return gl.LUMINANCE;
 		if ( p === LuminanceAlphaFormat ) return gl.LUMINANCE_ALPHA;

examples/webgpu_storage_buffer.html

@@ -26,7 +26,7 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { storageObject, vec3, uv, uint, float, tslFn, instanceIndex, MeshBasicNodeMaterial } from 'three/nodes';
+			import { storageObject, If, 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';
@@ -47,28 +47,42 @@
 
 				// texture
 
-				const typeSize = 1; // 1:'float', 2:'vec2', 4:'vec4' -> use power of 2
-				const size = 1024;
+				const size = 1024; // non power of two buffer size is not well supported in WebGPU
 
-				const array = new Array( size * typeSize ).fill( 0 );
+				const type = [ 'float', 'vec2', 'vec3', 'vec4' ];
 
-				const type = [ 'float', 'vec2', 'vec3', 'vec4' ][ typeSize - 1 ];
 
-				const arrayBuffer = new StorageBufferAttribute( new Float32Array( array ), typeSize );
+				const arrayBufferNodes = [];
 
-				const arrayBufferNode = storageObject( arrayBuffer, type, size );
+				for ( let i = 0; i < type.length; i ++ ) {
 
+					const typeSize = i + 1;
+					const array = new Array( size * typeSize ).fill( 0 );
+
+					const arrayBuffer = new StorageBufferAttribute( new Float32Array( array ), typeSize );
+
+					arrayBufferNodes.push( storageObject( arrayBuffer, type[ i ], size ) );
+
+				}
 
 				const computeInitOrder = tslFn( () => {
 
-					arrayBufferNode.element( instanceIndex ).assign( uint( instanceIndex.div( typeSize ) ) );
+					for ( let i = 0; i < type.length; i ++ ) {
+
+						arrayBufferNodes[ i ].element( instanceIndex ).assign( instanceIndex );
+
+					}
 
 				} );
 
 				const computeInvertOrder = tslFn( () => {
 
-					const invertIndex = arrayBufferNode.element( float( size ).sub( instanceIndex ) );
-					arrayBufferNode.element( instanceIndex ).assign( invertIndex );
+					for ( let i = 0; i < type.length; i ++ ) {
+
+						const invertIndex = arrayBufferNodes[ i ].element( uint( size ).sub( instanceIndex ) );
+						arrayBufferNodes[ i ].element( instanceIndex ).assign( invertIndex );
+
+					}
 
 				} );
 
@@ -82,14 +96,54 @@
 
 				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 );
+					const index = uint( uv().x.mul( size ).floor() ).toVar();
+
+					If( index.greaterThanEqual( size ), () => index.assign( uint( size ).sub( 1 ) ) );
+
+					const color = vec3( 0, 0, 0 ).toVar();
+
+					If( uv().y.greaterThan( 0.0 ), () => {
+
+						const indexValue = arrayBufferNodes[ 0 ].element( index ).toVar();
+						const value = float( indexValue ).div( float( size ) ).mul( 20 ).floor().div( 20 );
+			
+						color.assign( vec3( value, 0, 0 ) );
+
+					} );
 
-					return vec3( value, value, value );
+					If( uv().y.greaterThan( 0.25 ), () => {
+
+						const indexValue = arrayBufferNodes[ 1 ].element( index ).toVar();
+						const value = float( indexValue ).div( float( size ) ).mul( 20 ).floor().div( 20 );
+			
+						color.assign( vec3( 0, value, 0 ) );
+
+					} );
+
+					If( uv().y.greaterThan( 0.5 ), () => {
+
+						const indexValue = arrayBufferNodes[ 2 ].element( index ).toVar();
+						const value = float( indexValue ).div( float( size ) ).mul( 20 ).floor().div( 20 );
+			
+						color.assign( vec3( 0, 0, value ) );
+
+					} );
+
+					If( uv().y.greaterThan( 0.75 ), () => {
+
+						const indexValue = arrayBufferNodes[ 3 ].element( index ).toVar();
+						const value = float( indexValue ).div( float( size ) ).mul( 20 ).floor().div( 20 );
+			
+						color.assign( vec3( value, value, value ) );
+
+					} );
+
+					return color;
 
 				} )();
 			
+				// TODO: Add toAttribute() test
+
 				//
 
 				const plane = new THREE.Mesh( new THREE.PlaneGeometry( 1, 1 ), material );