Skip to content

TSL: Ensure memory alignment for struct() #31151

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
May 23, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 0 additions & 17 deletions src/extras/DataUtils.js
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
5 changes: 2 additions & 3 deletions src/nodes/accessors/Arrays.js
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down
42 changes: 42 additions & 0 deletions src/nodes/core/NodeUtils.js
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand Down
28 changes: 24 additions & 4 deletions src/nodes/core/StructTypeNode.js
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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;

}

Expand Down
32 changes: 14 additions & 18 deletions src/renderers/common/UniformsGroup.js
Original file line number Diff line number Diff line change
Expand Up @@ -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;

}

Expand Down
5 changes: 3 additions & 2 deletions src/renderers/webgpu/utils/WebGPUAttributeUtils.js
Original file line number Diff line number Diff line change
@@ -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' ]],
Expand Down Expand Up @@ -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,
Expand Down