Skip to content

Commit c9bb800

Browse files
sunagsunag
authored
TSL: Ensure memory alignment for struct() (#31151)
* improve alignToBoundary * cleanup * adds more robust verification * add `getMemoryLengthFromType()` * new approach * remove `DataUtils.alignToBoundary()` * cleanup * updates
1 parent 8301152 commit c9bb800

File tree

6 files changed

+85
-44
lines changed

6 files changed

+85
-44
lines changed

‎src/extras/DataUtils.js‎

Lines changed: 0 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -207,23 +207,6 @@ class DataUtils {
207207

208208
}
209209

210-
/**
211-
* Aligns a given byte length to the nearest 4-byte boundary.
212-
*
213-
* This function ensures that the returned byte length is a multiple of 4,
214-
* which is often required for memory alignment in certain systems or formats.
215-
*
216-
* @param {number} byteLength - The original byte length to align.
217-
* @returns {number} The aligned byte length, which is a multiple of 4.
218-
*/
219-
static alignTo4ByteBoundary( byteLength ) {
220-
221-
// ensure 4 byte alignment, see #20441
222-
223-
return byteLength + ( ( 4 - ( byteLength % 4 ) ) % 4 );
224-
225-
}
226-
227210
}
228211

229212
export {

‎src/nodes/accessors/Arrays.js‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ import StorageInstancedBufferAttribute from '../../renderers/common/StorageInsta
22
import StorageBufferAttribute from '../../renderers/common/StorageBufferAttribute.js';
33
import { storage } from './StorageBufferNode.js';
44
import { getLengthFromType, getTypedArrayFromType } from '../core/NodeUtils.js';
5-
import { DataUtils } from '../../extras/DataUtils.js';
65

76
/**
87
* TSL function for creating a storage buffer node with a configured `StorageBufferAttribute`.
@@ -19,7 +18,7 @@ export const attributeArray = ( count, type = 'float' ) => {
1918

2019
if ( type.isStruct === true ) {
2120

22-
itemSize = DataUtils.alignTo4ByteBoundary( type.layout.getLength() );
21+
itemSize = type.layout.getLength();
2322
typedArray = getTypedArrayFromType( 'float' );
2423

2524
} else {
@@ -51,7 +50,7 @@ export const instancedArray = ( count, type = 'float' ) => {
5150

5251
if ( type.isStruct === true ) {
5352

54-
itemSize = DataUtils.alignTo4ByteBoundary( type.layout.getLength() );
53+
itemSize = type.layout.getLength();
5554
typedArray = getTypedArrayFromType( 'float' );
5655

5756
} else {

‎src/nodes/core/NodeUtils.js‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -236,6 +236,48 @@ export function getLengthFromType( type ) {
236236

237237
}
238238

239+
/**
240+
* Returns the gpu memory length for the given data type.
241+
*
242+
* @method
243+
* @param {string} type - The data type.
244+
* @return {number} The length.
245+
*/
246+
export function getMemoryLengthFromType( type ) {
247+
248+
if ( /float|int|uint/.test( type ) ) return 1;
249+
if ( /vec2/.test( type ) ) return 2;
250+
if ( /vec3/.test( type ) ) return 3;
251+
if ( /vec4/.test( type ) ) return 4;
252+
if ( /mat2/.test( type ) ) return 4;
253+
if ( /mat3/.test( type ) ) return 12;
254+
if ( /mat4/.test( type ) ) return 16;
255+
256+
console.error( 'THREE.TSL: Unsupported type:', type );
257+
258+
}
259+
260+
/**
261+
* Returns the byte boundary for the given data type.
262+
*
263+
* @method
264+
* @param {string} type - The data type.
265+
* @return {number} The byte boundary.
266+
*/
267+
export function getByteBoundaryFromType( type ) {
268+
269+
if ( /float|int|uint/.test( type ) ) return 4;
270+
if ( /vec2/.test( type ) ) return 8;
271+
if ( /vec3/.test( type ) ) return 16;
272+
if ( /vec4/.test( type ) ) return 16;
273+
if ( /mat2/.test( type ) ) return 16;
274+
if ( /mat3/.test( type ) ) return 48;
275+
if ( /mat4/.test( type ) ) return 64;
276+
277+
console.error( 'THREE.TSL: Unsupported type:', type );
278+
279+
}
280+
239281
/**
240282
* Returns the data type for the given value.
241283
*

‎src/nodes/core/StructTypeNode.js‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11

22
import Node from './Node.js';
3-
import { getLengthFromType } from './NodeUtils.js';
3+
import { getByteBoundaryFromType, getMemoryLengthFromType } from './NodeUtils.js';
4+
import { GPU_CHUNK_BYTES } from '../../renderers/common/Constants.js';
45

56
/**
67
* Generates a layout for struct members.
@@ -86,15 +87,34 @@ class StructTypeNode extends Node {
8687
*/
8788
getLength() {
8889

89-
let length = 0;
90+
let offset = 0; // global buffer offset in bytes
9091

9192
for ( const member of this.membersLayout ) {
9293

93-
length += getLengthFromType( member.type );
94+
const type = member.type;
95+
96+
const itemSize = getMemoryLengthFromType( type ) * Float32Array.BYTES_PER_ELEMENT;
97+
const boundary = getByteBoundaryFromType( type );
98+
99+
const chunkOffset = offset % GPU_CHUNK_BYTES; // offset in the current chunk
100+
const chunkPadding = chunkOffset % boundary; // required padding to match boundary
101+
const chunkStart = chunkOffset + chunkPadding; // start position in the current chunk for the data
102+
103+
offset += chunkPadding;
104+
105+
// Check for chunk overflow
106+
if ( chunkStart !== 0 && ( GPU_CHUNK_BYTES - chunkStart ) < itemSize ) {
107+
108+
// Add padding to the end of the chunk
109+
offset += ( GPU_CHUNK_BYTES - chunkStart );
110+
111+
}
112+
113+
offset += itemSize;
94114

95115
}
96116

97-
return length;
117+
return ( Math.ceil( offset / GPU_CHUNK_BYTES ) * GPU_CHUNK_BYTES ) / Float32Array.BYTES_PER_ELEMENT;
98118

99119
}
100120

‎src/renderers/common/UniformsGroup.js‎

Lines changed: 14 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -129,38 +129,34 @@ class UniformsGroup extends UniformBuffer {
129129
*/
130130
get byteLength() {
131131

132+
const bytesPerElement = this.bytesPerElement;
133+
132134
let offset = 0; // global buffer offset in bytes
133135

134136
for ( let i = 0, l = this.uniforms.length; i < l; i ++ ) {
135137

136138
const uniform = this.uniforms[ i ];
137139

138-
const { boundary, itemSize } = uniform;
139-
140-
// offset within a single chunk in bytes
141-
142-
const chunkOffset = offset % GPU_CHUNK_BYTES;
143-
const remainingSizeInChunk = GPU_CHUNK_BYTES - chunkOffset;
144-
145-
// conformance tests
146-
147-
if ( chunkOffset !== 0 && ( remainingSizeInChunk - boundary ) < 0 ) {
148-
149-
// check for chunk overflow
140+
const boundary = uniform.boundary;
141+
const itemSize = uniform.itemSize * bytesPerElement; // size of the uniform in bytes
150142

151-
offset += ( GPU_CHUNK_BYTES - chunkOffset );
143+
const chunkOffset = offset % GPU_CHUNK_BYTES; // offset in the current chunk
144+
const chunkPadding = chunkOffset % boundary; // required padding to match boundary
145+
const chunkStart = chunkOffset + chunkPadding; // start position in the current chunk for the data
152146

153-
} else if ( chunkOffset % boundary !== 0 ) {
147+
offset += chunkPadding;
154148

155-
// check for correct alignment
149+
// Check for chunk overflow
150+
if ( chunkStart !== 0 && ( GPU_CHUNK_BYTES - chunkStart ) < itemSize ) {
156151

157-
offset += ( chunkOffset % boundary );
152+
// Add padding to the end of the chunk
153+
offset += ( GPU_CHUNK_BYTES - chunkStart );
158154

159155
}
160156

161-
uniform.offset = ( offset / this.bytesPerElement );
157+
uniform.offset = offset / bytesPerElement;
162158

163-
offset += ( itemSize * this.bytesPerElement );
159+
offset += itemSize;
164160

165161
}
166162

‎src/renderers/webgpu/utils/WebGPUAttributeUtils.js‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import { GPUInputStepMode } from './WebGPUConstants.js';
22

33
import { Float16BufferAttribute } from '../../../core/BufferAttribute.js';
4-
import { DataUtils } from '../../../extras/DataUtils.js';
54

65
const typedArraysToVertexFormatPrefix = new Map( [
76
[ Int8Array, [ 'sint8', 'snorm8' ]],
@@ -114,7 +113,9 @@ class WebGPUAttributeUtils {
114113

115114
}
116115

117-
const size = DataUtils.alignTo4ByteBoundary( array.byteLength );
116+
// ensure 4 byte alignment
117+
const byteLength = array.byteLength;
118+
const size = byteLength + ( ( 4 - ( byteLength % 4 ) ) % 4 );
118119

119120
buffer = device.createBuffer( {
120121
label: bufferAttribute.name,

0 commit comments

Comments
 (0)