TSL: Ensure memory alignment for `struct()` (#31151)

* improve alignToBoundary

* cleanup

* adds more robust verification

* add `getMemoryLengthFromType()`

* new approach

* remove `DataUtils.alignToBoundary()`

* cleanup

* updates

src/extras/DataUtils.js

@@ -207,23 +207,6 @@ class DataUtils {
 
 	}
 
-	/**
-	 * Aligns a given byte length to the nearest 4-byte boundary.
-	 *
-	 * This function ensures that the returned byte length is a multiple of 4,
-	 * which is often required for memory alignment in certain systems or formats.
-	 *
-	 * @param {number} byteLength - The original byte length to align.
-	 * @returns {number} The aligned byte length, which is a multiple of 4.
-	 */
-	static alignTo4ByteBoundary( byteLength ) {
-
-		// ensure 4 byte alignment, see #20441
-
-		return byteLength + ( ( 4 - ( byteLength % 4 ) ) % 4 );
-
-	}
-
 }
 
 export {

src/nodes/accessors/Arrays.js

@@ -2,7 +2,6 @@ import StorageInstancedBufferAttribute from '../../renderers/common/StorageInsta
 import StorageBufferAttribute from '../../renderers/common/StorageBufferAttribute.js';
 import { storage } from './StorageBufferNode.js';
 import { getLengthFromType, getTypedArrayFromType } from '../core/NodeUtils.js';
-import { DataUtils } from '../../extras/DataUtils.js';
 
 /**
  * TSL function for creating a storage buffer node with a configured `StorageBufferAttribute`.
@@ -19,7 +18,7 @@ export const attributeArray = ( count, type = 'float' ) => {
 
 	if ( type.isStruct === true ) {
 
-		itemSize = DataUtils.alignTo4ByteBoundary( type.layout.getLength() );
+		itemSize = type.layout.getLength();
 		typedArray = getTypedArrayFromType( 'float' );
 
 	} else {
@@ -51,7 +50,7 @@ export const instancedArray = ( count, type = 'float' ) => {
 
 	if ( type.isStruct === true ) {
 
-		itemSize = DataUtils.alignTo4ByteBoundary( type.layout.getLength() );
+		itemSize = type.layout.getLength();
 		typedArray = getTypedArrayFromType( 'float' );
 
 	} else {

src/nodes/core/NodeUtils.js

@@ -236,6 +236,48 @@ export function getLengthFromType( type ) {
 
 }
 
+/**
+ * Returns the gpu memory length for the given data type.
+ *
+ * @method
+ * @param {string} type - The data type.
+ * @return {number} The length.
+ */
+export function getMemoryLengthFromType( type ) {
+
+	if ( /float|int|uint/.test( type ) ) return 1;
+	if ( /vec2/.test( type ) ) return 2;
+	if ( /vec3/.test( type ) ) return 3;
+	if ( /vec4/.test( type ) ) return 4;
+	if ( /mat2/.test( type ) ) return 4;
+	if ( /mat3/.test( type ) ) return 12;
+	if ( /mat4/.test( type ) ) return 16;
+
+	console.error( 'THREE.TSL: Unsupported type:', type );
+
+}
+
+/**
+ * Returns the byte boundary for the given data type.
+ *
+ * @method
+ * @param {string} type - The data type.
+ * @return {number} The byte boundary.
+ */
+export function getByteBoundaryFromType( type ) {
+
+	if ( /float|int|uint/.test( type ) ) return 4;
+	if ( /vec2/.test( type ) ) return 8;
+	if ( /vec3/.test( type ) ) return 16;
+	if ( /vec4/.test( type ) ) return 16;
+	if ( /mat2/.test( type ) ) return 16;
+	if ( /mat3/.test( type ) ) return 48;
+	if ( /mat4/.test( type ) ) return 64;
+
+	console.error( 'THREE.TSL: Unsupported type:', type );
+
+}
+
 /**
  * Returns the data type for the given value.
  *

src/nodes/core/StructTypeNode.js

@@ -1,6 +1,7 @@
 
 import Node from './Node.js';
-import { getLengthFromType } from './NodeUtils.js';
+import { getByteBoundaryFromType, getMemoryLengthFromType } from './NodeUtils.js';
+import { GPU_CHUNK_BYTES } from '../../renderers/common/Constants.js';
 
 /**
  * Generates a layout for struct members.
@@ -86,15 +87,34 @@ class StructTypeNode extends Node {
 	 */
 	getLength() {
 
-		let length = 0;
+		let offset = 0; // global buffer offset in bytes
 
 		for ( const member of this.membersLayout ) {
 
-			length += getLengthFromType( member.type );
+			const type = member.type;
+
+			const itemSize = getMemoryLengthFromType( type ) * Float32Array.BYTES_PER_ELEMENT;
+			const boundary = getByteBoundaryFromType( type );
+
+			const chunkOffset = offset % GPU_CHUNK_BYTES; // offset in the current chunk
+			const chunkPadding = chunkOffset % boundary; // required padding to match boundary
+			const chunkStart = chunkOffset + chunkPadding; // start position in the current chunk for the data
+
+			offset += chunkPadding;
+
+			// Check for chunk overflow
+			if ( chunkStart !== 0 && ( GPU_CHUNK_BYTES - chunkStart ) < itemSize ) {
+
+				// Add padding to the end of the chunk
+				offset += ( GPU_CHUNK_BYTES - chunkStart );
+
+			}
+
+			offset += itemSize;
 
 		}
 
-		return length;
+		return ( Math.ceil( offset / GPU_CHUNK_BYTES ) * GPU_CHUNK_BYTES ) / Float32Array.BYTES_PER_ELEMENT;
 
 	}
 

src/renderers/common/UniformsGroup.js

@@ -129,38 +129,34 @@ class UniformsGroup extends UniformBuffer {
 	 */
 	get byteLength() {
 
+		const bytesPerElement = this.bytesPerElement;
+
 		let offset = 0; // global buffer offset in bytes
 
 		for ( let i = 0, l = this.uniforms.length; i < l; i ++ ) {
 
 			const uniform = this.uniforms[ i ];
 
-			const { boundary, itemSize } = uniform;
-
-			// offset within a single chunk in bytes
-
-			const chunkOffset = offset % GPU_CHUNK_BYTES;
-			const remainingSizeInChunk = GPU_CHUNK_BYTES - chunkOffset;
-
-			// conformance tests
-
-			if ( chunkOffset !== 0 && ( remainingSizeInChunk - boundary ) < 0 ) {
-
-				// check for chunk overflow
+			const boundary = uniform.boundary;
+			const itemSize = uniform.itemSize * bytesPerElement; // size of the uniform in bytes
 
-				offset += ( GPU_CHUNK_BYTES - chunkOffset );
+			const chunkOffset = offset % GPU_CHUNK_BYTES; // offset in the current chunk
+			const chunkPadding = chunkOffset % boundary; // required padding to match boundary
+			const chunkStart = chunkOffset + chunkPadding; // start position in the current chunk for the data
 
-			} else if ( chunkOffset % boundary !== 0 ) {
+			offset += chunkPadding;
 
-				// check for correct alignment
+			// Check for chunk overflow
+			if ( chunkStart !== 0 && ( GPU_CHUNK_BYTES - chunkStart ) < itemSize ) {
 
-				offset += ( chunkOffset % boundary );
+				// Add padding to the end of the chunk
+				offset += ( GPU_CHUNK_BYTES - chunkStart );
 
 			}
 
-			uniform.offset = ( offset / this.bytesPerElement );
+			uniform.offset = offset / bytesPerElement;
 
-			offset += ( itemSize * this.bytesPerElement );
+			offset += itemSize;
 
 		}
 

src/renderers/webgpu/utils/WebGPUAttributeUtils.js

@@ -1,7 +1,6 @@
 import { GPUInputStepMode } from './WebGPUConstants.js';
 
 import { Float16BufferAttribute } from '../../../core/BufferAttribute.js';
-import { DataUtils } from '../../../extras/DataUtils.js';
 
 const typedArraysToVertexFormatPrefix = new Map( [
 	[ Int8Array, [ 'sint8', 'snorm8' ]],
@@ -114,7 +113,9 @@ class WebGPUAttributeUtils {
 
 			}
 
-			const size = DataUtils.alignTo4ByteBoundary( array.byteLength );
+			// ensure 4 byte alignment
+			const byteLength = array.byteLength;
+			const size = byteLength + ( ( 4 - ( byteLength % 4 ) ) % 4 );
 
 			buffer = device.createBuffer( {
 				label: bufferAttribute.name,