From 6b3783919e4743763bd0fe578072e630ad418f21 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 11:37:21 +0100 Subject: [PATCH 01/14] impr: add SoA write path --- packages/typegpu/src/core/buffer/buffer.ts | 23 ++++++- packages/typegpu/src/data/soaIO.ts | 74 ++++++++++++++++++++++ packages/typegpu/src/data/wgslTypes.ts | 17 ++++- packages/typegpu/tests/buffer.test.ts | 70 ++++++++++++++++++++ 4 files changed, 180 insertions(+), 4 deletions(-) create mode 100644 packages/typegpu/src/data/soaIO.ts diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index 21ca3956cc..2d28e93c32 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -5,7 +5,8 @@ import type { AnyData } from '../../data/dataTypes.ts'; import { getWriteInstructions } from '../../data/partialIO.ts'; import { sizeOf } from '../../data/sizeOf.ts'; import type { BaseData } from '../../data/wgslTypes.ts'; -import { isWgslArray, isWgslData } from '../../data/wgslTypes.ts'; +import { writeSoA } from '../../data/soaIO.ts'; +import { isWgslArray, isWgslData, isWgslStruct } from '../../data/wgslTypes.ts'; import type { StorageFlag } from '../../extension.ts'; import type { TgpuNamable } from '../../shared/meta.ts'; import { getName, setName } from '../../shared/meta.ts'; @@ -322,6 +323,26 @@ class TgpuBufferImpl implements TgpuBuffer { return; } + // SoA path: struct-of-arrays input for array-of-structs schema + if ( + isWgslArray(this.dataType) && + isWgslStruct(this.dataType.elementType) && + !Array.isArray(data) && + typeof data === 'object' && + data !== null + ) { + const firstValue = Object.values(data as Record)[0]; + if (firstValue !== undefined && ArrayBuffer.isView(firstValue)) { + writeSoA( + new Uint8Array(target), + this.dataType, + data as Record, + startOffset, + ); + return; + } + } + const dataView = new DataView(target); const isLittleEndian = endianness === 'little'; diff --git a/packages/typegpu/src/data/soaIO.ts b/packages/typegpu/src/data/soaIO.ts new file mode 100644 index 0000000000..582a445ecb --- /dev/null +++ b/packages/typegpu/src/data/soaIO.ts @@ -0,0 +1,74 @@ +import { invariant } from '../errors.ts'; +import { roundUp } from '../mathUtils.ts'; +import { alignmentOf } from './alignmentOf.ts'; +import { offsetsForProps } from './offsets.ts'; +import { sizeOf } from './sizeOf.ts'; +import type { WgslArray, WgslStruct } from './wgslTypes.ts'; +import { isMat, isMat3x3f } from './wgslTypes.ts'; + +/** + * Writes struct-of-arrays (SoA) data into a GPU-layout (AoS) target buffer. + * + * Each key in `soaData` is a struct field name mapped to a packed TypedArray + * containing that field's values for all elements (no inter-element padding). + * This function scatters those packed arrays into the correctly padded AoS layout. + */ +export function writeSoA( + target: Uint8Array, + arraySchema: WgslArray, + soaData: Record, + startOffset: number, +): void { + const structSchema = arraySchema.elementType as WgslStruct; + const offsets = offsetsForProps(structSchema); + const elementStride = roundUp(sizeOf(structSchema), alignmentOf(structSchema)); + const elementCount = arraySchema.elementCount; + + for (const key in structSchema.propTypes) { + const fieldSchema = structSchema.propTypes[key]; + if (fieldSchema === undefined) { + continue; + } + const srcArray = soaData[key]; + if (srcArray === undefined) { + continue; + } + + const fieldOffset = offsets[key]?.offset; + invariant(fieldOffset !== undefined, `Field ${key} not found in struct schema`); + const srcBytes = new Uint8Array(srcArray.buffer, srcArray.byteOffset, srcArray.byteLength); + + if (isMat(fieldSchema)) { + // Matrices may have internal column padding (mat3x3f: columns are vec3f + // stored at 16-byte stride, but packed data has 12 bytes per column). + const dim = isMat3x3f(fieldSchema) ? 3 : fieldSchema.type === 'mat2x2f' ? 2 : 4; + const compSize = 4; // all current matrix types use f32 + const packedColumnSize = dim * compSize; + const gpuColumnStride = roundUp(packedColumnSize, alignmentOf(fieldSchema)); + const packedElementSize = dim * packedColumnSize; // total packed bytes per matrix + + for (let i = 0; i < elementCount; i++) { + const dstBase = startOffset + i * elementStride + fieldOffset; + const srcBase = i * packedElementSize; + for (let col = 0; col < dim; col++) { + target.set( + srcBytes.subarray( + srcBase + col * packedColumnSize, + srcBase + col * packedColumnSize + packedColumnSize, + ), + dstBase + col * gpuColumnStride, + ); + } + } + } else { + // Scalars and vectors: packed size equals sizeOf(field), no internal padding. + const fieldSize = sizeOf(fieldSchema); + for (let i = 0; i < elementCount; i++) { + target.set( + srcBytes.subarray(i * fieldSize, i * fieldSize + fieldSize), + startOffset + i * elementStride + fieldOffset, + ); + } + } + } +} diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts index 7756a776f2..a9cf8a0cda 100644 --- a/packages/typegpu/src/data/wgslTypes.ts +++ b/packages/typegpu/src/data/wgslTypes.ts @@ -53,9 +53,9 @@ export interface NumberArrayView { } /** - * Maps a scalar or vector element schema to the corresponding TypedArray type. + * Maps a scalar, vector, or matrix element schema to the corresponding TypedArray type. */ -export type TypedArrayFor = T extends Vec2f | Vec3f | Vec4f | F32 +export type TypedArrayFor = T extends Vec2f | Vec3f | Vec4f | F32 | Mat2x2f | Mat3x3f | Mat4x4f ? Float32Array : T extends Vec2h | Vec3h | Vec4h | F16 ? Float16Array @@ -67,6 +67,14 @@ export type TypedArrayFor = T extends Vec2f | Vec3f | Vec4f | F32 ? Uint16Array : never; +/** + * Maps struct properties to a record of TypedArrays (Struct-of-Arrays input format). + * If any property resolves to `never` (e.g. nested structs), the type becomes unconstructable. + */ +export type SoAInputFor> = { + [K in keyof TProps]: TypedArrayFor; +}; + /** * Vector infix notation. * @@ -1167,7 +1175,10 @@ export interface WgslArray extends Bas // Type-tokens, not available at runtime readonly [$repr]: Infer[]; - readonly [$inRepr]: InferInput[] | TypedArrayFor; + readonly [$inRepr]: + | InferInput[] + | TypedArrayFor + | (TElement extends WgslStruct ? SoAInputFor : never); readonly [$gpuRepr]: InferGPU[]; readonly [$reprPartial]: { idx: number; value: InferPartial }[] | undefined; readonly [$memIdent]: WgslArray>; diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index 848288c045..dc0d59ff48 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -978,4 +978,74 @@ describe('ValidateBufferSchema', () => { createMyBuffer(d.f32, ['uniform']); createMyBuffer(d.unorm8x4, ['vertex']); }); + + it('should write struct-of-arrays (SoA) data to an array-of-structs buffer', ({ + root, + device, + }) => { + const Particle = d.struct({ + pos: d.vec3f, + vel: d.f32, + }); + // vec3f alignment=16, size=12; f32 alignment=4, size=4 + // struct layout: pos @ 0 (12 bytes), vel @ 12 (4 bytes) → struct size=16, alignment=16 + // element stride = 16 + + const schema = d.arrayOf(Particle, 2); + const buffer = root.createBuffer(schema); + const rawBuffer = root.unwrap(buffer); + + buffer.write({ + pos: new Float32Array([1, 2, 3, 4, 5, 6]), + vel: new Float32Array([10, 20]), + }); + + const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; + const result = new Float32Array(uploadedBuffer); + + // Element 0: pos=(1,2,3) @ offset 0, vel=10 @ offset 12 + // Element 1: pos=(4,5,6) @ offset 16, vel=20 @ offset 28 + expect([...result]).toStrictEqual([ + 1, + 2, + 3, + 10, // element 0 + 4, + 5, + 6, + 20, // element 1 + ]); + }); + + it('should write SoA data with integer fields', ({ root, device }) => { + const Entry = d.struct({ + id: d.u32, + heading: d.vec3i, + }); + // u32: align=4, size=4; vec3i: align=16, size=12 + // struct layout: id @ 0 (4 bytes), [12 bytes padding], heading @ 16 (12 bytes) → struct size=32 + + const schema = d.arrayOf(Entry, 2); + const buffer = root.createBuffer(schema); + root.unwrap(buffer); + + buffer.write({ + id: new Uint32Array([100, 200]), + heading: new Int32Array([1, 2, 3, 4, 5, 6]), + }); + + const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; + + // Element 0: id=100 @ byte 0, heading=(1,2,3) @ byte 16 + // Element 1: id=200 @ byte 32, heading=(4,5,6) @ byte 48 + const ids = [ + new DataView(uploadedBuffer).getUint32(0, true), + new DataView(uploadedBuffer).getUint32(32, true), + ]; + const headings = [new Int32Array(uploadedBuffer, 16, 3), new Int32Array(uploadedBuffer, 48, 3)]; + + expect(ids).toStrictEqual([100, 200]); + expect([...headings[0]!]).toStrictEqual([1, 2, 3]); + expect([...headings[1]!]).toStrictEqual([4, 5, 6]); + }); }); From 29eb88a1a9ea92aca176f29069fdaa5be61f9d8c Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 12:28:21 +0100 Subject: [PATCH 02/14] MORE --- packages/typegpu/src/core/buffer/buffer.ts | 19 +- packages/typegpu/src/data/index.ts | 1 + packages/typegpu/src/data/soaIO.ts | 187 ++++++++++++--- packages/typegpu/src/data/wgslTypes.ts | 26 ++- packages/typegpu/tests/buffer.test.ts | 258 +++++++++++++++++++++ 5 files changed, 442 insertions(+), 49 deletions(-) diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index 2d28e93c32..e3e4e5c1d2 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -5,7 +5,7 @@ import type { AnyData } from '../../data/dataTypes.ts'; import { getWriteInstructions } from '../../data/partialIO.ts'; import { sizeOf } from '../../data/sizeOf.ts'; import type { BaseData } from '../../data/wgslTypes.ts'; -import { writeSoA } from '../../data/soaIO.ts'; +import { getSoANaturalSize, isSoACompatibleField, writeSoA } from '../../data/soaIO.ts'; import { isWgslArray, isWgslData, isWgslStruct } from '../../data/wgslTypes.ts'; import type { StorageFlag } from '../../extension.ts'; import type { TgpuNamable } from '../../shared/meta.ts'; @@ -323,7 +323,7 @@ class TgpuBufferImpl implements TgpuBuffer { return; } - // SoA path: struct-of-arrays input for array-of-structs schema + // SoA path if ( isWgslArray(this.dataType) && isWgslStruct(this.dataType.elementType) && @@ -331,13 +331,20 @@ class TgpuBufferImpl implements TgpuBuffer { typeof data === 'object' && data !== null ) { - const firstValue = Object.values(data as Record)[0]; - if (firstValue !== undefined && ArrayBuffer.isView(firstValue)) { + const soaData = data as Record; + const values = Object.values(soaData); + const isSoAInput = + values.length > 0 && + values.every(ArrayBuffer.isView) && + Object.values(this.dataType.elementType.propTypes).every(isSoACompatibleField); + + if (isSoAInput) { writeSoA( new Uint8Array(target), this.dataType, - data as Record, + soaData as Record, startOffset, + endOffset, ); return; } @@ -380,7 +387,7 @@ class TgpuBufferImpl implements TgpuBuffer { roundUp(sizeOf(this.dataType.elementType), alignmentOf(this.dataType.elementType)) : ArrayBuffer.isView(data) || data instanceof ArrayBuffer ? data.byteLength - : undefined; + : getSoANaturalSize(this.dataType, data); const naturalEndOffset = naturalSize !== undefined ? Math.min(startOffset + naturalSize, bufferSize) : undefined; diff --git a/packages/typegpu/src/data/index.ts b/packages/typegpu/src/data/index.ts index 7b51b0467c..7cd58ca2ea 100644 --- a/packages/typegpu/src/data/index.ts +++ b/packages/typegpu/src/data/index.ts @@ -89,6 +89,7 @@ export type { Mat4x4f, matBase, Ptr, + SoAInputFor, Size, StorableData, U16, diff --git a/packages/typegpu/src/data/soaIO.ts b/packages/typegpu/src/data/soaIO.ts index 582a445ecb..1fc77bd2ad 100644 --- a/packages/typegpu/src/data/soaIO.ts +++ b/packages/typegpu/src/data/soaIO.ts @@ -3,8 +3,146 @@ import { roundUp } from '../mathUtils.ts'; import { alignmentOf } from './alignmentOf.ts'; import { offsetsForProps } from './offsets.ts'; import { sizeOf } from './sizeOf.ts'; -import type { WgslArray, WgslStruct } from './wgslTypes.ts'; -import { isMat, isMat3x3f } from './wgslTypes.ts'; +import type { BaseData, WgslArray, WgslStruct } from './wgslTypes.ts'; +import { isMat, isMat2x2f, isMat3x3f, isWgslArray, isWgslStruct } from './wgslTypes.ts'; + +function getPackedMatrixLayout(schema: BaseData) { + if (!isMat(schema)) { + return undefined; + } + + const dim = isMat3x3f(schema) ? 3 : isMat2x2f(schema) ? 2 : 4; + const packedColumnSize = dim * 4; + + return { + dim, + packedColumnSize, + packedSize: dim * packedColumnSize, + } as const; +} + +export function packedSizeOf(schema: BaseData): number { + const matrixLayout = getPackedMatrixLayout(schema); + if (matrixLayout) { + return matrixLayout.packedSize; + } + + if (isWgslArray(schema)) { + return schema.elementCount * packedSizeOf(schema.elementType); + } + + return sizeOf(schema); +} + +export function inferSoAElementCount( + arraySchema: WgslArray, + soaData: Record, +): number | undefined { + const structSchema = arraySchema.elementType as WgslStruct; + let inferredCount: number | undefined; + + for (const key in soaData) { + const srcArray = soaData[key]; + const fieldSchema = structSchema.propTypes[key]; + if (srcArray === undefined || fieldSchema === undefined) { + continue; + } + + const fieldPackedSize = packedSizeOf(fieldSchema); + if (fieldPackedSize === 0) { + continue; + } + + const fieldElementCount = Math.floor(srcArray.byteLength / fieldPackedSize); + inferredCount = + inferredCount === undefined ? fieldElementCount : Math.min(inferredCount, fieldElementCount); + } + + return inferredCount; +} + +export function isSoACompatibleField(schema: BaseData): boolean { + if (isWgslArray(schema)) { + return isSoACompatibleField(schema.elementType); + } + + return !isWgslStruct(schema); +} + +export function getSoANaturalSize(dataType: BaseData, data: unknown): number | undefined { + if ( + !isWgslArray(dataType) || + !isWgslStruct(dataType.elementType) || + Array.isArray(data) || + typeof data !== 'object' || + data === null + ) { + return undefined; + } + + const soaData = data as Record; + const values = Object.values(soaData); + const isSoAInput = + values.length > 0 && + values.every(ArrayBuffer.isView) && + Object.values(dataType.elementType.propTypes).every(isSoACompatibleField); + + if (!isSoAInput) { + return undefined; + } + + const elementCount = inferSoAElementCount(dataType, soaData as Record); + if (elementCount === undefined) { + return undefined; + } + + const elementStride = roundUp(sizeOf(dataType.elementType), alignmentOf(dataType.elementType)); + return elementCount * elementStride; +} + +function writePackedValue( + target: Uint8Array, + schema: BaseData, + srcBytes: Uint8Array, + dstOffset: number, + srcOffset: number, +): void { + const matrixLayout = getPackedMatrixLayout(schema); + if (matrixLayout) { + const gpuColumnStride = roundUp(matrixLayout.packedColumnSize, alignmentOf(schema)); + + for (let col = 0; col < matrixLayout.dim; col++) { + target.set( + srcBytes.subarray( + srcOffset + col * matrixLayout.packedColumnSize, + srcOffset + col * matrixLayout.packedColumnSize + matrixLayout.packedColumnSize, + ), + dstOffset + col * gpuColumnStride, + ); + } + + return; + } + + if (isWgslArray(schema)) { + const packedElementSize = packedSizeOf(schema.elementType); + const gpuElementStride = roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType)); + + for (let i = 0; i < schema.elementCount; i++) { + writePackedValue( + target, + schema.elementType, + srcBytes, + dstOffset + i * gpuElementStride, + srcOffset + i * packedElementSize, + ); + } + + return; + } + + target.set(srcBytes.subarray(srcOffset, srcOffset + sizeOf(schema)), dstOffset); +} /** * Writes struct-of-arrays (SoA) data into a GPU-layout (AoS) target buffer. @@ -18,11 +156,14 @@ export function writeSoA( arraySchema: WgslArray, soaData: Record, startOffset: number, + endOffset: number, ): void { const structSchema = arraySchema.elementType as WgslStruct; const offsets = offsetsForProps(structSchema); const elementStride = roundUp(sizeOf(structSchema), alignmentOf(structSchema)); - const elementCount = arraySchema.elementCount; + const startElement = Math.floor(startOffset / elementStride); + const endElement = Math.min(arraySchema.elementCount, Math.ceil(endOffset / elementStride)); + const elementCount = Math.max(0, endElement - startElement); for (const key in structSchema.propTypes) { const fieldSchema = structSchema.propTypes[key]; @@ -38,37 +179,15 @@ export function writeSoA( invariant(fieldOffset !== undefined, `Field ${key} not found in struct schema`); const srcBytes = new Uint8Array(srcArray.buffer, srcArray.byteOffset, srcArray.byteLength); - if (isMat(fieldSchema)) { - // Matrices may have internal column padding (mat3x3f: columns are vec3f - // stored at 16-byte stride, but packed data has 12 bytes per column). - const dim = isMat3x3f(fieldSchema) ? 3 : fieldSchema.type === 'mat2x2f' ? 2 : 4; - const compSize = 4; // all current matrix types use f32 - const packedColumnSize = dim * compSize; - const gpuColumnStride = roundUp(packedColumnSize, alignmentOf(fieldSchema)); - const packedElementSize = dim * packedColumnSize; // total packed bytes per matrix - - for (let i = 0; i < elementCount; i++) { - const dstBase = startOffset + i * elementStride + fieldOffset; - const srcBase = i * packedElementSize; - for (let col = 0; col < dim; col++) { - target.set( - srcBytes.subarray( - srcBase + col * packedColumnSize, - srcBase + col * packedColumnSize + packedColumnSize, - ), - dstBase + col * gpuColumnStride, - ); - } - } - } else { - // Scalars and vectors: packed size equals sizeOf(field), no internal padding. - const fieldSize = sizeOf(fieldSchema); - for (let i = 0; i < elementCount; i++) { - target.set( - srcBytes.subarray(i * fieldSize, i * fieldSize + fieldSize), - startOffset + i * elementStride + fieldOffset, - ); - } + const packedFieldSize = packedSizeOf(fieldSchema); + for (let i = 0; i < elementCount; i++) { + writePackedValue( + target, + fieldSchema, + srcBytes, + (startElement + i) * elementStride + fieldOffset, + i * packedFieldSize, + ); } } } diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts index a9cf8a0cda..bb0d354d30 100644 --- a/packages/typegpu/src/data/wgslTypes.ts +++ b/packages/typegpu/src/data/wgslTypes.ts @@ -55,25 +55,33 @@ export interface NumberArrayView { /** * Maps a scalar, vector, or matrix element schema to the corresponding TypedArray type. */ -export type TypedArrayFor = T extends Vec2f | Vec3f | Vec4f | F32 | Mat2x2f | Mat3x3f | Mat4x4f +export type TypedArrayFor = T extends F32 | Vec2f | Vec3f | Vec4f | Mat2x2f | Mat3x3f | Mat4x4f ? Float32Array - : T extends Vec2h | Vec3h | Vec4h | F16 + : T extends F16 | Vec2h | Vec3h | Vec4h ? Float16Array - : T extends Vec2i | Vec3i | Vec4i | I32 + : T extends I32 | Vec2i | Vec3i | Vec4i ? Int32Array - : T extends Vec2u | Vec3u | Vec4u | U32 + : T extends U32 | Vec2u | Vec3u | Vec4u ? Uint32Array : T extends U16 ? Uint16Array : never; +type UnwrapWgslArray = T extends WgslArray ? UnwrapWgslArray : T; +type PackedSoAInputFor = TypedArrayFor>; + +type SoAFieldsFor> = { + [K in keyof T as [PackedSoAInputFor] extends [never] ? never : K]: PackedSoAInputFor; +}; + /** - * Maps struct properties to a record of TypedArrays (Struct-of-Arrays input format). - * If any property resolves to `never` (e.g. nested structs), the type becomes unconstructable. + * Maps struct properties to a record of TypedArrays (Struct-of-Arrays input format) - if not possible, resolves to never. */ -export type SoAInputFor> = { - [K in keyof TProps]: TypedArrayFor; -}; +export type SoAInputFor> = [keyof T] extends [ + keyof SoAFieldsFor, +] + ? Prettify> + : never; /** * Vector infix notation. diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index dc0d59ff48..96b5dc93c6 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -1,6 +1,7 @@ import { attest } from '@ark/attest'; import { describe, expect, expectTypeOf, vi } from 'vitest'; import * as d from '../src/data/index.ts'; +import type { SoAInputFor } from '../src/data/index.ts'; import { sizeOf } from '../src/data/sizeOf.ts'; import type { ValidateBufferSchema, ValidUsagesFor } from '../src/index.js'; import { getName } from '../src/shared/meta.ts'; @@ -1048,4 +1049,261 @@ describe('ValidateBufferSchema', () => { expect([...headings[0]!]).toStrictEqual([1, 2, 3]); expect([...headings[1]!]).toStrictEqual([4, 5, 6]); }); + + it('should accept SoA input for struct fields that are fixed-size arrays of primitives', () => { + type Test = { + a: d.F32; + b: d.Vec3u; + c: d.Mat4x4f; + d: d.WgslArray; + e: d.WgslArray; + f: d.WgslArray; + }; + + expectTypeOf>().toEqualTypeOf<{ + a: Float32Array; + b: Uint32Array; + c: Float32Array; + d: Float32Array; + e: Int32Array; + f: Float32Array; + }>(); + }); + + it('should reject SoA input for struct fields that contain nested structs', () => { + const Nested = d.struct({ + x: d.f32, + }); + + type Test = { + a: d.F32; + nested: typeof Nested; + }; + + expectTypeOf>().toEqualTypeOf(); + }); + + it('should write SoA data for struct fields that are fixed-size arrays of primitives', ({ + root, + device, + }) => { + const Entry = d.struct({ + id: d.u32, + values: d.arrayOf(d.f32, 3), + }); + + const schema = d.arrayOf(Entry, 2); + const buffer = root.createBuffer(schema); + root.unwrap(buffer); + + buffer.write({ + id: new Uint32Array([10, 20]), + values: new Float32Array([1, 2, 3, 4, 5, 6]), + }); + + const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; + const ids = [ + new DataView(uploadedBuffer).getUint32(0, true), + new DataView(uploadedBuffer).getUint32(16, true), + ]; + const values = [ + new Float32Array(uploadedBuffer, 4, 3), + new Float32Array(uploadedBuffer, 20, 3), + ]; + + expect(ids).toStrictEqual([10, 20]); + expect([...values[0]!]).toStrictEqual([1, 2, 3]); + expect([...values[1]!]).toStrictEqual([4, 5, 6]); + }); + + it('should write SoA data for struct fields that are arrays of padded vectors', ({ + root, + device, + }) => { + const Entry = d.struct({ + values: d.arrayOf(d.vec3f, 2), + }); + + const schema = d.arrayOf(Entry, 2); + const buffer = root.createBuffer(schema); + root.unwrap(buffer); + + buffer.write({ + values: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]), + }); + + const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; + const result = new Float32Array(uploadedBuffer); + + expect([...result]).toStrictEqual([1, 2, 3, 0, 4, 5, 6, 0, 7, 8, 9, 0, 10, 11, 12, 0]); + }); + + it('should write SoA data for struct fields that are arrays of padded matrices', ({ + root, + device, + }) => { + const Entry = d.struct({ + basis: d.arrayOf(d.mat3x3f, 2), + }); + + const schema = d.arrayOf(Entry, 1); + const buffer = root.createBuffer(schema); + root.unwrap(buffer); + + buffer.write({ + basis: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]), + }); + + const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; + const result = new Float32Array(uploadedBuffer); + + expect([...result]).toStrictEqual([ + 1, 2, 3, 0, 4, 5, 6, 0, 7, 8, 9, 0, 10, 11, 12, 0, 13, 14, 15, 0, 16, 17, 18, 0, + ]); + }); + + it('should write SoA data only for the middle two elements when using startOffset and endOffset', ({ + root, + device, + }) => { + const Entry = d.struct({ + id: d.u32, + values: d.arrayOf(d.vec3f, 2), + }); + + const schema = d.arrayOf(Entry, 4); + const buffer = root.createBuffer(schema); + const rawBuffer = root.unwrap(buffer); + const startLayout = d.memoryLayoutOf(schema, (a) => a[1]); + const endLayout = d.memoryLayoutOf(schema, (a) => a[3]); + const idLayout = d.memoryLayoutOf(Entry, (e) => e.id); + const value0Layout = d.memoryLayoutOf(Entry, (e) => e.values[0]); + const value1Layout = d.memoryLayoutOf(Entry, (e) => e.values[1]); + const stride = sizeOf(Entry); + + buffer.write( + { + id: new Uint32Array([30, 40]), + values: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]), + }, + { + startOffset: startLayout.offset, + endOffset: endLayout.offset, + }, + ); + + expect(device.mock.queue.writeBuffer.mock.calls).toStrictEqual([ + [ + rawBuffer, + startLayout.offset, + expect.any(ArrayBuffer), + startLayout.offset, + endLayout.offset - startLayout.offset, + ], + ]); + + const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; + const firstBase = startLayout.offset; + const secondBase = startLayout.offset + stride; + const ids = [ + new DataView(uploadedBuffer).getUint32(firstBase + idLayout.offset, true), + new DataView(uploadedBuffer).getUint32(secondBase + idLayout.offset, true), + ]; + const valueVectors = [ + new Float32Array(uploadedBuffer, firstBase + value0Layout.offset, 3), + new Float32Array(uploadedBuffer, firstBase + value1Layout.offset, 3), + new Float32Array(uploadedBuffer, secondBase + value0Layout.offset, 3), + new Float32Array(uploadedBuffer, secondBase + value1Layout.offset, 3), + ]; + const untouchedPrefix = new Uint32Array(uploadedBuffer, 0, startLayout.offset / 4); + const untouchedSuffix = new Uint32Array( + uploadedBuffer, + endLayout.offset, + (sizeOf(schema) - endLayout.offset) / 4, + ); + + expect([...untouchedPrefix]).toStrictEqual( + Array.from({ length: startLayout.offset / 4 }, () => 0), + ); + expect(ids).toStrictEqual([30, 40]); + expect([...valueVectors[0]!]).toStrictEqual([1, 2, 3]); + expect([...valueVectors[1]!]).toStrictEqual([4, 5, 6]); + expect([...valueVectors[2]!]).toStrictEqual([7, 8, 9]); + expect([...valueVectors[3]!]).toStrictEqual([10, 11, 12]); + expect([...untouchedSuffix]).toStrictEqual( + Array.from({ length: (sizeOf(schema) - endLayout.offset) / 4 }, () => 0), + ); + }); + + it('should infer the SoA write range from provided data when endOffset is omitted', ({ + root, + device, + }) => { + const Entry = d.struct({ + id: d.u32, + values: d.arrayOf(d.vec3f, 2), + }); + + const schema = d.arrayOf(Entry, 4); + const buffer = root.createBuffer(schema); + const rawBuffer = root.unwrap(buffer); + const startLayout = d.memoryLayoutOf(schema, (a) => a[1]); + const endLayout = d.memoryLayoutOf(schema, (a) => a[3]); + const idLayout = d.memoryLayoutOf(Entry, (e) => e.id); + const value0Layout = d.memoryLayoutOf(Entry, (e) => e.values[0]); + const value1Layout = d.memoryLayoutOf(Entry, (e) => e.values[1]); + const stride = sizeOf(Entry); + + buffer.write( + { + id: new Uint32Array([30, 40]), + values: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]), + }, + { + startOffset: startLayout.offset, + }, + ); + + expect(device.mock.queue.writeBuffer.mock.calls).toStrictEqual([ + [ + rawBuffer, + startLayout.offset, + expect.any(ArrayBuffer), + startLayout.offset, + endLayout.offset - startLayout.offset, + ], + ]); + + const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; + const firstBase = startLayout.offset; + const secondBase = startLayout.offset + stride; + const ids = [ + new DataView(uploadedBuffer).getUint32(firstBase + idLayout.offset, true), + new DataView(uploadedBuffer).getUint32(secondBase + idLayout.offset, true), + ]; + const valueVectors = [ + new Float32Array(uploadedBuffer, firstBase + value0Layout.offset, 3), + new Float32Array(uploadedBuffer, firstBase + value1Layout.offset, 3), + new Float32Array(uploadedBuffer, secondBase + value0Layout.offset, 3), + new Float32Array(uploadedBuffer, secondBase + value1Layout.offset, 3), + ]; + const untouchedPrefix = new Uint32Array(uploadedBuffer, 0, startLayout.offset / 4); + const untouchedSuffix = new Uint32Array( + uploadedBuffer, + endLayout.offset, + (sizeOf(schema) - endLayout.offset) / 4, + ); + + expect([...untouchedPrefix]).toStrictEqual( + Array.from({ length: startLayout.offset / 4 }, () => 0), + ); + expect(ids).toStrictEqual([30, 40]); + expect([...valueVectors[0]!]).toStrictEqual([1, 2, 3]); + expect([...valueVectors[1]!]).toStrictEqual([4, 5, 6]); + expect([...valueVectors[2]!]).toStrictEqual([7, 8, 9]); + expect([...valueVectors[3]!]).toStrictEqual([10, 11, 12]); + expect([...untouchedSuffix]).toStrictEqual( + Array.from({ length: (sizeOf(schema) - endLayout.offset) / 4 }, () => 0), + ); + }); }); From 0b6766ab3cd9518d33245cc5d5a47180bc8b12f8 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 13:51:19 +0100 Subject: [PATCH 03/14] NEW API JUST DROPPED --- packages/typegpu/src/common/index.ts | 1 + packages/typegpu/src/common/writeSoA.ts | 27 ++++++ packages/typegpu/src/core/buffer/buffer.ts | 93 +++++++++++---------- packages/typegpu/src/core/root/init.ts | 10 +-- packages/typegpu/src/core/root/rootTypes.ts | 11 ++- packages/typegpu/src/data/soaIO.ts | 2 +- packages/typegpu/src/data/wgslTypes.ts | 5 +- packages/typegpu/tests/buffer.test.ts | 93 +++++++++++++++++++-- 8 files changed, 175 insertions(+), 67 deletions(-) create mode 100644 packages/typegpu/src/common/writeSoA.ts diff --git a/packages/typegpu/src/common/index.ts b/packages/typegpu/src/common/index.ts index 607f3438fa..9249008035 100644 --- a/packages/typegpu/src/common/index.ts +++ b/packages/typegpu/src/common/index.ts @@ -1,3 +1,4 @@ // NOTE: This is a barrel file, internal files should not import things from this file export { fullScreenTriangle } from './fullScreenTriangle.ts'; +export { writeSoA } from './writeSoA.ts'; diff --git a/packages/typegpu/src/common/writeSoA.ts b/packages/typegpu/src/common/writeSoA.ts new file mode 100644 index 0000000000..628fbf04b7 --- /dev/null +++ b/packages/typegpu/src/common/writeSoA.ts @@ -0,0 +1,27 @@ +import { getSoANaturalSize, writeSoA as scatterSoA } from '../data/soaIO.ts'; +import { sizeOf } from '../data/sizeOf.ts'; +import type { BaseData, SoAInputFor, WgslArray, WgslStruct } from '../data/wgslTypes.ts'; +import type { BufferWriteOptions, TgpuBuffer } from '../core/buffer/buffer.ts'; + +export function writeSoA>( + buffer: TgpuBuffer>>, + data: SoAInputFor, + options?: BufferWriteOptions, +): void { + const arrayBuffer = buffer.arrayBuffer; + const startOffset = options?.startOffset ?? 0; + const bufferSize = sizeOf(buffer.dataType); + const naturalSize = getSoANaturalSize(buffer.dataType, data); + const endOffset = + options?.endOffset ?? + (naturalSize === undefined ? bufferSize : Math.min(startOffset + naturalSize, bufferSize)); + + scatterSoA( + new Uint8Array(arrayBuffer), + buffer.dataType, + data as Record, + startOffset, + endOffset, + ); + buffer.write(arrayBuffer, { startOffset, endOffset }); +} diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index e3e4e5c1d2..66ba895096 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -5,8 +5,7 @@ import type { AnyData } from '../../data/dataTypes.ts'; import { getWriteInstructions } from '../../data/partialIO.ts'; import { sizeOf } from '../../data/sizeOf.ts'; import type { BaseData } from '../../data/wgslTypes.ts'; -import { getSoANaturalSize, isSoACompatibleField, writeSoA } from '../../data/soaIO.ts'; -import { isWgslArray, isWgslData, isWgslStruct } from '../../data/wgslTypes.ts'; +import { isWgslArray, isWgslData } from '../../data/wgslTypes.ts'; import type { StorageFlag } from '../../extension.ts'; import type { TgpuNamable } from '../../shared/meta.ts'; import { getName, setName } from '../../shared/meta.ts'; @@ -112,11 +111,17 @@ export type BufferWriteOptions = { endOffset?: number; }; +export type BufferInitCallback = (buffer: TgpuBuffer) => void; +export type BufferInitialData = + | InferInput + | BufferInitCallback; + export interface TgpuBuffer extends TgpuNamable { readonly [$internal]: true; readonly resourceType: 'buffer'; readonly dataType: TData; readonly initial?: InferInput | undefined; + readonly arrayBuffer: ArrayBuffer; readonly buffer: GPUBuffer; readonly destroyed: boolean; @@ -153,7 +158,7 @@ export interface TgpuBuffer extends TgpuNamable { export function INTERNAL_createBuffer( group: ExperimentalTgpuRoot, typeSchema: TData, - initialOrBuffer?: InferInput | GPUBuffer, + initialOrBuffer?: BufferInitialData | GPUBuffer, ): TgpuBuffer { if (!isWgslData(typeSchema)) { return new TgpuBufferImpl(group, typeSchema, initialOrBuffer, ['storage', 'uniform']); @@ -193,6 +198,7 @@ class TgpuBufferImpl implements TgpuBuffer { private _ownBuffer: boolean; private _destroyed = false; private _hostBuffer: ArrayBuffer | undefined; + private _initialCallback: BufferInitCallback | undefined; readonly initial: InferInput | undefined; @@ -205,7 +211,7 @@ class TgpuBufferImpl implements TgpuBuffer { constructor( root: ExperimentalTgpuRoot, public readonly dataType: TData, - public readonly initialOrBuffer?: InferInput | GPUBuffer, + public readonly initialOrBuffer?: BufferInitialData | GPUBuffer, private readonly _disallowedUsages?: UsageLiteral[], ) { this.#device = root.device; @@ -214,7 +220,11 @@ class TgpuBufferImpl implements TgpuBuffer { this._buffer = initialOrBuffer; } else { this._ownBuffer = true; - this.initial = initialOrBuffer; + if (typeof initialOrBuffer === 'function') { + this._initialCallback = initialOrBuffer as BufferInitCallback; + } else { + this.initial = initialOrBuffer; + } } } @@ -227,12 +237,16 @@ class TgpuBufferImpl implements TgpuBuffer { this._buffer = this.#device.createBuffer({ size: sizeOf(this.dataType), usage: this.flags, - mappedAtCreation: !!this.initial, + mappedAtCreation: !!this.initial || !!this._initialCallback, label: getName(this) ?? '', }); - if (this.initial) { - this._writeToTarget(this._buffer.getMappedRange(), this.initial); + if (this.initial || this._initialCallback) { + if (this._initialCallback) { + this._initialCallback(this); + } else if (this.initial) { + this._writeToTarget(this._buffer.getMappedRange(), this.initial); + } this._buffer.unmap(); } } @@ -244,6 +258,19 @@ class TgpuBufferImpl implements TgpuBuffer { return this._destroyed; } + get arrayBuffer(): ArrayBuffer { + const gpuBuffer = this.buffer; + if (gpuBuffer.mapState === 'mapped') { + return gpuBuffer.getMappedRange(); + } + + if (!this._hostBuffer) { + this._hostBuffer = new ArrayBuffer(sizeOf(this.dataType)); + } + + return this._hostBuffer; + } + $name(label: string) { setName(this, label); if (this._buffer) { @@ -323,33 +350,6 @@ class TgpuBufferImpl implements TgpuBuffer { return; } - // SoA path - if ( - isWgslArray(this.dataType) && - isWgslStruct(this.dataType.elementType) && - !Array.isArray(data) && - typeof data === 'object' && - data !== null - ) { - const soaData = data as Record; - const values = Object.values(soaData); - const isSoAInput = - values.length > 0 && - values.every(ArrayBuffer.isView) && - Object.values(this.dataType.elementType.propTypes).every(isSoACompatibleField); - - if (isSoAInput) { - writeSoA( - new Uint8Array(target), - this.dataType, - soaData as Record, - startOffset, - endOffset, - ); - return; - } - } - const dataView = new DataView(target); const isLittleEndian = endianness === 'little'; @@ -381,13 +381,15 @@ class TgpuBufferImpl implements TgpuBuffer { const bufferSize = sizeOf(this.dataType); const startOffset = options?.startOffset ?? 0; - const naturalSize = - isWgslArray(this.dataType) && Array.isArray(data) - ? data.length * - roundUp(sizeOf(this.dataType.elementType), alignmentOf(this.dataType.elementType)) - : ArrayBuffer.isView(data) || data instanceof ArrayBuffer - ? data.byteLength - : getSoANaturalSize(this.dataType, data); + let naturalSize: number | undefined = undefined; + if (isWgslArray(this.dataType) && Array.isArray(data)) { + const arrayData = data as unknown[]; + naturalSize = + arrayData.length * + roundUp(sizeOf(this.dataType.elementType), alignmentOf(this.dataType.elementType)); + } else if (ArrayBuffer.isView(data) || data instanceof ArrayBuffer) { + naturalSize = data.byteLength; + } const naturalEndOffset = naturalSize !== undefined ? Math.min(startOffset + naturalSize, bufferSize) : undefined; @@ -396,6 +398,9 @@ class TgpuBufferImpl implements TgpuBuffer { if (gpuBuffer.mapState === 'mapped') { const mapped = gpuBuffer.getMappedRange(); + if (data instanceof ArrayBuffer && data === mapped) { + return; + } this._writeToTarget(mapped, data, options); return; } @@ -404,7 +409,9 @@ class TgpuBufferImpl implements TgpuBuffer { this._hostBuffer = new ArrayBuffer(bufferSize); } - this._writeToTarget(this._hostBuffer, data, options); + if (!(data instanceof ArrayBuffer && data === this._hostBuffer)) { + this._writeToTarget(this._hostBuffer, data, options); + } this.#device.queue.writeBuffer(gpuBuffer, startOffset, this._hostBuffer, startOffset, size); } diff --git a/packages/typegpu/src/core/root/init.ts b/packages/typegpu/src/core/root/init.ts index c874341760..8fad9a5951 100644 --- a/packages/typegpu/src/core/root/init.ts +++ b/packages/typegpu/src/core/root/init.ts @@ -9,7 +9,7 @@ import type { AnyWgslData, BaseData, v3u, Vec3u, WgslArray } from '../../data/wg import { invariant } from '../../errors.ts'; import { WeakMemo } from '../../memo.ts'; import { clearTextureUtilsCache } from '../texture/textureUtils.ts'; -import type { InferInput } from '../../shared/repr.ts'; +import type { BufferInitialData } from '../buffer/buffer.ts'; import { $getNameForward, $internal } from '../../shared/symbols.ts'; import type { ExtractBindGroupInputFromLayout, @@ -439,14 +439,14 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu createBuffer( typeSchema: TData, - initialOrBuffer?: InferInput | GPUBuffer, + initialOrBuffer?: BufferInitialData | GPUBuffer, ): TgpuBuffer { return INTERNAL_createBuffer(this, typeSchema, initialOrBuffer); } createUniform( typeSchema: TData, - initialOrBuffer?: InferInput | GPUBuffer, + initialOrBuffer?: BufferInitialData | GPUBuffer, ): TgpuUniform { const buffer = INTERNAL_createBuffer(this, typeSchema, initialOrBuffer) // oxlint-disable-next-line typescript/no-explicit-any -- i'm sure it's fine @@ -457,7 +457,7 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu createMutable( typeSchema: TData, - initialOrBuffer?: InferInput | GPUBuffer, + initialOrBuffer?: BufferInitialData | GPUBuffer, ): TgpuMutable { const buffer = INTERNAL_createBuffer(this, typeSchema, initialOrBuffer) // oxlint-disable-next-line typescript/no-explicit-any -- i'm sure it's fine @@ -468,7 +468,7 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu createReadonly( typeSchema: TData, - initialOrBuffer?: InferInput | GPUBuffer, + initialOrBuffer?: BufferInitialData | GPUBuffer, ): TgpuReadonly { const buffer = INTERNAL_createBuffer(this, typeSchema, initialOrBuffer) // oxlint-disable-next-line typescript/no-explicit-any -- i'm sure it's fine diff --git a/packages/typegpu/src/core/root/rootTypes.ts b/packages/typegpu/src/core/root/rootTypes.ts index 2df3c32525..f7735efc89 100644 --- a/packages/typegpu/src/core/root/rootTypes.ts +++ b/packages/typegpu/src/core/root/rootTypes.ts @@ -15,7 +15,6 @@ import type { import type { ExtractInvalidSchemaError, InferGPURecord, - InferInput, IsValidBufferSchema, IsValidStorageSchema, IsValidUniformSchema, @@ -31,7 +30,7 @@ import type { import type { LogGeneratorOptions } from '../../tgsl/consoleLog/types.ts'; import type { ShaderGenerator } from '../../tgsl/shaderGenerator.ts'; import type { Unwrapper } from '../../unwrapper.ts'; -import type { TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; +import type { BufferInitialData, TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; import type { TgpuMutable, TgpuReadonly, TgpuUniform } from '../buffer/bufferShorthand.ts'; import type { TgpuFixedComparisonSampler, TgpuFixedSampler } from '../sampler/sampler.ts'; import type { IORecord } from '../function/fnTypes.ts'; @@ -836,7 +835,7 @@ export interface TgpuRoot extends Unwrapper, WithBinding { createBuffer( typeSchema: ValidateBufferSchema, // NoInfer is there to infer the schema type just based on the first parameter - initial?: InferInput>, + initial?: BufferInitialData>, ): TgpuBuffer; /** @@ -864,7 +863,7 @@ export interface TgpuRoot extends Unwrapper, WithBinding { createUniform( typeSchema: ValidateUniformSchema, // NoInfer is there to infer the schema type just based on the first parameter - initial?: InferInput>, + initial?: BufferInitialData>, ): TgpuUniform; /** @@ -891,7 +890,7 @@ export interface TgpuRoot extends Unwrapper, WithBinding { createMutable( typeSchema: ValidateStorageSchema, // NoInfer is there to infer the schema type just based on the first parameter - initial?: InferInput>, + initial?: BufferInitialData>, ): TgpuMutable; /** @@ -918,7 +917,7 @@ export interface TgpuRoot extends Unwrapper, WithBinding { createReadonly( typeSchema: ValidateStorageSchema, // NoInfer is there to infer the schema type just based on the first parameter - initial?: InferInput>, + initial?: BufferInitialData>, ): TgpuReadonly; /** diff --git a/packages/typegpu/src/data/soaIO.ts b/packages/typegpu/src/data/soaIO.ts index 1fc77bd2ad..738934c846 100644 --- a/packages/typegpu/src/data/soaIO.ts +++ b/packages/typegpu/src/data/soaIO.ts @@ -84,7 +84,7 @@ export function getSoANaturalSize(dataType: BaseData, data: unknown): number | u const values = Object.values(soaData); const isSoAInput = values.length > 0 && - values.every(ArrayBuffer.isView) && + values.every((value) => ArrayBuffer.isView(value)) && Object.values(dataType.elementType.propTypes).every(isSoACompatibleField); if (!isSoAInput) { diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts index bb0d354d30..a6738e7452 100644 --- a/packages/typegpu/src/data/wgslTypes.ts +++ b/packages/typegpu/src/data/wgslTypes.ts @@ -1183,10 +1183,7 @@ export interface WgslArray extends Bas // Type-tokens, not available at runtime readonly [$repr]: Infer[]; - readonly [$inRepr]: - | InferInput[] - | TypedArrayFor - | (TElement extends WgslStruct ? SoAInputFor : never); + readonly [$inRepr]: InferInput[] | TypedArrayFor; readonly [$gpuRepr]: InferGPU[]; readonly [$reprPartial]: { idx: number; value: InferPartial }[] | undefined; readonly [$memIdent]: WgslArray>; diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index 96b5dc93c6..9b6f8f5ad5 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -1,7 +1,8 @@ import { attest } from '@ark/attest'; import { describe, expect, expectTypeOf, vi } from 'vitest'; +import * as common from '../src/common/index.ts'; import * as d from '../src/data/index.ts'; -import type { SoAInputFor } from '../src/data/index.ts'; +import type { InferInput, SoAInputFor } from '../src/data/index.ts'; import { sizeOf } from '../src/data/sizeOf.ts'; import type { ValidateBufferSchema, ValidUsagesFor } from '../src/index.js'; import { getName } from '../src/shared/meta.ts'; @@ -101,6 +102,48 @@ describe('TgpuBuffer', () => { expect(root.device.queue.writeBuffer).toBeCalledWith(mockBuffer, 0, new ArrayBuffer(64), 0, 64); }); + it('should initialize a buffer from a mapped callback using common.writeSoA', ({ root }) => { + const Entry = d.struct({ + id: d.u32, + values: d.arrayOf(d.vec3f, 2), + }); + + const buffer = root.createBuffer(d.arrayOf(Entry, 2), (mappedBuffer) => { + common.writeSoA(mappedBuffer, { + id: new Uint32Array([10, 20]), + values: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]), + }); + }); + + const rawBuffer = root.unwrap(buffer); + const writtenBuffer = vi.mocked(rawBuffer.getMappedRange).mock.results[0]?.value as ArrayBuffer; + + expect(root.device.createBuffer).toBeCalledWith( + expect.objectContaining({ + mappedAtCreation: true, + size: sizeOf(d.arrayOf(Entry, 2)), + usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, + }), + ); + expect(root.device.queue.writeBuffer).not.toHaveBeenCalled(); + const ids = [ + new DataView(writtenBuffer).getUint32(0, true), + new DataView(writtenBuffer).getUint32(48, true), + ]; + const values = [ + new Float32Array(writtenBuffer, 16, 3), + new Float32Array(writtenBuffer, 32, 3), + new Float32Array(writtenBuffer, 64, 3), + new Float32Array(writtenBuffer, 80, 3), + ]; + + expect(ids).toStrictEqual([10, 20]); + expect([...values[0]!]).toStrictEqual([1, 2, 3]); + expect([...values[1]!]).toStrictEqual([4, 5, 6]); + expect([...values[2]!]).toStrictEqual([7, 8, 9]); + expect([...values[3]!]).toStrictEqual([10, 11, 12]); + }); + it('should write to a mapped buffer', ({ root }) => { const mappedBuffer = root.device.createBuffer({ size: 12, @@ -996,7 +1039,7 @@ describe('ValidateBufferSchema', () => { const buffer = root.createBuffer(schema); const rawBuffer = root.unwrap(buffer); - buffer.write({ + common.writeSoA(buffer, { pos: new Float32Array([1, 2, 3, 4, 5, 6]), vel: new Float32Array([10, 20]), }); @@ -1030,7 +1073,7 @@ describe('ValidateBufferSchema', () => { const buffer = root.createBuffer(schema); root.unwrap(buffer); - buffer.write({ + common.writeSoA(buffer, { id: new Uint32Array([100, 200]), heading: new Int32Array([1, 2, 3, 4, 5, 6]), }); @@ -1083,6 +1126,38 @@ describe('ValidateBufferSchema', () => { expectTypeOf>().toEqualTypeOf(); }); + it('should not accept SoA input through buffer.write', ({ root }) => { + const Entry = d.struct({ + id: d.u32, + values: d.arrayOf(d.f32, 3), + }); + const buffer = root.createBuffer(d.arrayOf(Entry, 2)); + + if (false) { + // @ts-expect-error SoA writes go through common.writeSoA, not buffer.write + buffer.write({ + id: new Uint32Array([10, 20]), + values: new Float32Array([1, 2, 3, 4, 5, 6]), + }); + } + }); + + it('should keep SoA input out of nested array fields in InferInput', () => { + const Child = d.struct({ + x: d.f32, + }); + const Parent = d.struct({ + children: d.arrayOf(Child, 2), + }); + const TopLevel = d.arrayOf(Parent, 4); + + expectTypeOf>().toEqualTypeOf<{ + children: { x: number }[]; + }>(); + + expectTypeOf>().toEqualTypeOf<{ children: { x: number }[] }[]>(); + }); + it('should write SoA data for struct fields that are fixed-size arrays of primitives', ({ root, device, @@ -1096,7 +1171,7 @@ describe('ValidateBufferSchema', () => { const buffer = root.createBuffer(schema); root.unwrap(buffer); - buffer.write({ + common.writeSoA(buffer, { id: new Uint32Array([10, 20]), values: new Float32Array([1, 2, 3, 4, 5, 6]), }); @@ -1128,7 +1203,7 @@ describe('ValidateBufferSchema', () => { const buffer = root.createBuffer(schema); root.unwrap(buffer); - buffer.write({ + common.writeSoA(buffer, { values: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]), }); @@ -1150,7 +1225,7 @@ describe('ValidateBufferSchema', () => { const buffer = root.createBuffer(schema); root.unwrap(buffer); - buffer.write({ + common.writeSoA(buffer, { basis: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]), }); @@ -1181,7 +1256,8 @@ describe('ValidateBufferSchema', () => { const value1Layout = d.memoryLayoutOf(Entry, (e) => e.values[1]); const stride = sizeOf(Entry); - buffer.write( + common.writeSoA( + buffer, { id: new Uint32Array([30, 40]), values: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]), @@ -1254,7 +1330,8 @@ describe('ValidateBufferSchema', () => { const value1Layout = d.memoryLayoutOf(Entry, (e) => e.values[1]); const stride = sizeOf(Entry); - buffer.write( + common.writeSoA( + buffer, { id: new Uint32Array([30, 40]), values: new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]), From de2430d16b32106e96e75de2916dafa6a671f4cb Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 14:16:19 +0100 Subject: [PATCH 04/14] cache mapped buffer --- packages/typegpu/src/core/buffer/buffer.ts | 38 +++++++++++++++++----- packages/typegpu/tests/buffer.test.ts | 1 + 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index 66ba895096..1914c7aa16 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -198,6 +198,7 @@ class TgpuBufferImpl implements TgpuBuffer { private _ownBuffer: boolean; private _destroyed = false; private _hostBuffer: ArrayBuffer | undefined; + private _mappedRange: ArrayBuffer | undefined; private _initialCallback: BufferInitCallback | undefined; readonly initial: InferInput | undefined; @@ -245,9 +246,9 @@ class TgpuBufferImpl implements TgpuBuffer { if (this._initialCallback) { this._initialCallback(this); } else if (this.initial) { - this._writeToTarget(this._buffer.getMappedRange(), this.initial); + this._writeToTarget(this._getMappedRange(), this.initial); } - this._buffer.unmap(); + this._unmapBuffer(); } } @@ -261,7 +262,7 @@ class TgpuBufferImpl implements TgpuBuffer { get arrayBuffer(): ArrayBuffer { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - return gpuBuffer.getMappedRange(); + return this._getMappedRange(); } if (!this._hostBuffer) { @@ -271,6 +272,24 @@ class TgpuBufferImpl implements TgpuBuffer { return this._hostBuffer; } + private _getMappedRange(): ArrayBuffer { + if (!this._buffer || this._buffer.mapState !== 'mapped') { + throw new Error('Buffer is not mapped.'); + } + + this._mappedRange ??= this._buffer.getMappedRange(); + return this._mappedRange; + } + + private _unmapBuffer(): void { + if (!this._buffer || this._buffer.mapState !== 'mapped') { + return; + } + + this._mappedRange = undefined; + this._buffer.unmap(); + } + $name(label: string) { setName(this, label); if (this._buffer) { @@ -397,7 +416,7 @@ class TgpuBufferImpl implements TgpuBuffer { const size = endOffset - startOffset; if (gpuBuffer.mapState === 'mapped') { - const mapped = gpuBuffer.getMappedRange(); + const mapped = this._getMappedRange(); if (data instanceof ArrayBuffer && data === mapped) { return; } @@ -421,7 +440,7 @@ class TgpuBufferImpl implements TgpuBuffer { const instructions = getWriteInstructions(this.dataType, data); if (gpuBuffer.mapState === 'mapped') { - const mappedRange = gpuBuffer.getMappedRange(); + const mappedRange = this._getMappedRange(); const mappedView = new Uint8Array(mappedRange); for (const instruction of instructions) { @@ -444,7 +463,7 @@ class TgpuBufferImpl implements TgpuBuffer { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - new Uint8Array(gpuBuffer.getMappedRange()).fill(0); + new Uint8Array(this._getMappedRange()).fill(0); return; } @@ -468,15 +487,15 @@ class TgpuBufferImpl implements TgpuBuffer { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - const mapped = gpuBuffer.getMappedRange(); + const mapped = this._getMappedRange(); return readData(new BufferReader(mapped), this.dataType); } if (gpuBuffer.usage & GPUBufferUsage.MAP_READ) { await gpuBuffer.mapAsync(GPUMapMode.READ); - const mapped = gpuBuffer.getMappedRange(); + const mapped = this._getMappedRange(); const res = readData(new BufferReader(mapped), this.dataType); - gpuBuffer.unmap(); + this._unmapBuffer(); return res; } @@ -508,6 +527,7 @@ class TgpuBufferImpl implements TgpuBuffer { return; } this._destroyed = true; + this._mappedRange = undefined; if (this._ownBuffer) { this._buffer?.destroy(); } diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index 9b6f8f5ad5..69a285b12e 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -125,6 +125,7 @@ describe('TgpuBuffer', () => { usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, }), ); + expect(rawBuffer.getMappedRange).toHaveBeenCalledTimes(1); expect(root.device.queue.writeBuffer).not.toHaveBeenCalled(); const ids = [ new DataView(writtenBuffer).getUint32(0, true), From b038ad431a48e08ee63d7f7245a6323efca0e1f2 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 14:23:42 +0100 Subject: [PATCH 05/14] trim outdated tests --- packages/typegpu/tests/buffer.test.ts | 53 +-------------------------- 1 file changed, 1 insertion(+), 52 deletions(-) diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index 69a285b12e..a5fff10ba0 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -1032,9 +1032,6 @@ describe('ValidateBufferSchema', () => { pos: d.vec3f, vel: d.f32, }); - // vec3f alignment=16, size=12; f32 alignment=4, size=4 - // struct layout: pos @ 0 (12 bytes), vel @ 12 (4 bytes) → struct size=16, alignment=16 - // element stride = 16 const schema = d.arrayOf(Particle, 2); const buffer = root.createBuffer(schema); @@ -1048,18 +1045,7 @@ describe('ValidateBufferSchema', () => { const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; const result = new Float32Array(uploadedBuffer); - // Element 0: pos=(1,2,3) @ offset 0, vel=10 @ offset 12 - // Element 1: pos=(4,5,6) @ offset 16, vel=20 @ offset 28 - expect([...result]).toStrictEqual([ - 1, - 2, - 3, - 10, // element 0 - 4, - 5, - 6, - 20, // element 1 - ]); + expect([...result]).toStrictEqual([1, 2, 3, 10, 4, 5, 6, 20]); }); it('should write SoA data with integer fields', ({ root, device }) => { @@ -1067,9 +1053,6 @@ describe('ValidateBufferSchema', () => { id: d.u32, heading: d.vec3i, }); - // u32: align=4, size=4; vec3i: align=16, size=12 - // struct layout: id @ 0 (4 bytes), [12 bytes padding], heading @ 16 (12 bytes) → struct size=32 - const schema = d.arrayOf(Entry, 2); const buffer = root.createBuffer(schema); root.unwrap(buffer); @@ -1081,8 +1064,6 @@ describe('ValidateBufferSchema', () => { const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; - // Element 0: id=100 @ byte 0, heading=(1,2,3) @ byte 16 - // Element 1: id=200 @ byte 32, heading=(4,5,6) @ byte 48 const ids = [ new DataView(uploadedBuffer).getUint32(0, true), new DataView(uploadedBuffer).getUint32(32, true), @@ -1127,38 +1108,6 @@ describe('ValidateBufferSchema', () => { expectTypeOf>().toEqualTypeOf(); }); - it('should not accept SoA input through buffer.write', ({ root }) => { - const Entry = d.struct({ - id: d.u32, - values: d.arrayOf(d.f32, 3), - }); - const buffer = root.createBuffer(d.arrayOf(Entry, 2)); - - if (false) { - // @ts-expect-error SoA writes go through common.writeSoA, not buffer.write - buffer.write({ - id: new Uint32Array([10, 20]), - values: new Float32Array([1, 2, 3, 4, 5, 6]), - }); - } - }); - - it('should keep SoA input out of nested array fields in InferInput', () => { - const Child = d.struct({ - x: d.f32, - }); - const Parent = d.struct({ - children: d.arrayOf(Child, 2), - }); - const TopLevel = d.arrayOf(Parent, 4); - - expectTypeOf>().toEqualTypeOf<{ - children: { x: number }[]; - }>(); - - expectTypeOf>().toEqualTypeOf<{ children: { x: number }[] }[]>(); - }); - it('should write SoA data for struct fields that are fixed-size arrays of primitives', ({ root, device, From 2f03d17e266895356423ac4892391e686f833b99 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 14:50:52 +0100 Subject: [PATCH 06/14] overloads === better type hints --- packages/typegpu/src/core/buffer/buffer.ts | 2 +- packages/typegpu/src/core/root/rootTypes.ts | 27 +++++++++++++++++---- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index 1914c7aa16..39228a9715 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -145,7 +145,7 @@ export interface TgpuBuffer extends TgpuNamable { as>(usage: T): UsageTypeToBufferUsage[T]; compileWriter(): void; - write(data: InferInput, options?: BufferWriteOptions): void; + write(data: Prettify>, options?: BufferWriteOptions): void; write(data: ArrayBuffer, options?: BufferWriteOptions): void; writePartial(data: InferPartial): void; clear(): void; diff --git a/packages/typegpu/src/core/root/rootTypes.ts b/packages/typegpu/src/core/root/rootTypes.ts index f7735efc89..05b9c3db32 100644 --- a/packages/typegpu/src/core/root/rootTypes.ts +++ b/packages/typegpu/src/core/root/rootTypes.ts @@ -15,6 +15,7 @@ import type { import type { ExtractInvalidSchemaError, InferGPURecord, + InferInput, IsValidBufferSchema, IsValidStorageSchema, IsValidUniformSchema, @@ -30,7 +31,7 @@ import type { import type { LogGeneratorOptions } from '../../tgsl/consoleLog/types.ts'; import type { ShaderGenerator } from '../../tgsl/shaderGenerator.ts'; import type { Unwrapper } from '../../unwrapper.ts'; -import type { BufferInitialData, TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; +import type { TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; import type { TgpuMutable, TgpuReadonly, TgpuUniform } from '../buffer/bufferShorthand.ts'; import type { TgpuFixedComparisonSampler, TgpuFixedSampler } from '../sampler/sampler.ts'; import type { IORecord } from '../function/fnTypes.ts'; @@ -832,10 +833,14 @@ export interface TgpuRoot extends Unwrapper, WithBinding { * @param typeSchema The type of data that this buffer will hold. * @param initial The initial value of the buffer. (optional) */ + createBuffer( + typeSchema: ValidateBufferSchema, + initializer: (buffer: TgpuBuffer) => void, + ): TgpuBuffer; createBuffer( typeSchema: ValidateBufferSchema, // NoInfer is there to infer the schema type just based on the first parameter - initial?: BufferInitialData>, + initial?: InferInput>, ): TgpuBuffer; /** @@ -860,10 +865,14 @@ export interface TgpuRoot extends Unwrapper, WithBinding { * @param typeSchema The type of data that this buffer will hold. * @param initial The initial value of the buffer. (optional) */ + createUniform( + typeSchema: ValidateUniformSchema, + initializer: (buffer: TgpuBuffer) => void, + ): TgpuUniform; createUniform( typeSchema: ValidateUniformSchema, // NoInfer is there to infer the schema type just based on the first parameter - initial?: BufferInitialData>, + initial?: InferInput>, ): TgpuUniform; /** @@ -887,10 +896,14 @@ export interface TgpuRoot extends Unwrapper, WithBinding { * @param typeSchema The type of data that this buffer will hold. * @param initial The initial value of the buffer. (optional) */ + createMutable( + typeSchema: ValidateStorageSchema, + initializer: (buffer: TgpuBuffer) => void, + ): TgpuMutable; createMutable( typeSchema: ValidateStorageSchema, // NoInfer is there to infer the schema type just based on the first parameter - initial?: BufferInitialData>, + initial?: InferInput>, ): TgpuMutable; /** @@ -914,10 +927,14 @@ export interface TgpuRoot extends Unwrapper, WithBinding { * @param typeSchema The type of data that this buffer will hold. * @param initial The initial value of the buffer. (optional) */ + createReadonly( + typeSchema: ValidateStorageSchema, + initializer: (buffer: TgpuBuffer) => void, + ): TgpuReadonly; createReadonly( typeSchema: ValidateStorageSchema, // NoInfer is there to infer the schema type just based on the first parameter - initial?: BufferInitialData>, + initial?: InferInput>, ): TgpuReadonly; /** From 9adcd5aa09ab6ed415ecb7993ae93cc1406e1e91 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 15:17:39 +0100 Subject: [PATCH 07/14] tweaks and comments --- packages/typegpu/src/common/writeSoA.ts | 7 +++-- packages/typegpu/src/core/buffer/buffer.ts | 6 +++- packages/typegpu/src/data/soaIO.ts | 35 ++++++++++++++++------ 3 files changed, 36 insertions(+), 12 deletions(-) diff --git a/packages/typegpu/src/common/writeSoA.ts b/packages/typegpu/src/common/writeSoA.ts index 628fbf04b7..368af6ab5f 100644 --- a/packages/typegpu/src/common/writeSoA.ts +++ b/packages/typegpu/src/common/writeSoA.ts @@ -1,4 +1,4 @@ -import { getSoANaturalSize, writeSoA as scatterSoA } from '../data/soaIO.ts'; +import { computeSoAByteLength, scatterSoA } from '../data/soaIO.ts'; import { sizeOf } from '../data/sizeOf.ts'; import type { BaseData, SoAInputFor, WgslArray, WgslStruct } from '../data/wgslTypes.ts'; import type { BufferWriteOptions, TgpuBuffer } from '../core/buffer/buffer.ts'; @@ -11,7 +11,10 @@ export function writeSoA>( const arrayBuffer = buffer.arrayBuffer; const startOffset = options?.startOffset ?? 0; const bufferSize = sizeOf(buffer.dataType); - const naturalSize = getSoANaturalSize(buffer.dataType, data); + const naturalSize = computeSoAByteLength( + buffer.dataType, + data as Record, + ); const endOffset = options?.endOffset ?? (naturalSize === undefined ? bufferSize : Math.min(startOffset + naturalSize, bufferSize)); diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index 39228a9715..bed001f91b 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -212,7 +212,7 @@ class TgpuBufferImpl implements TgpuBuffer { constructor( root: ExperimentalTgpuRoot, public readonly dataType: TData, - public readonly initialOrBuffer?: BufferInitialData | GPUBuffer, + initialOrBuffer?: BufferInitialData | GPUBuffer, private readonly _disallowedUsages?: UsageLiteral[], ) { this.#device = root.device; @@ -418,6 +418,8 @@ class TgpuBufferImpl implements TgpuBuffer { if (gpuBuffer.mapState === 'mapped') { const mapped = this._getMappedRange(); if (data instanceof ArrayBuffer && data === mapped) { + // The caller already wrote data directly into the mapped range + // via arrayBuffer. Nothing to do here return; } this._writeToTarget(mapped, data, options); @@ -428,6 +430,8 @@ class TgpuBufferImpl implements TgpuBuffer { this._hostBuffer = new ArrayBuffer(bufferSize); } + // If the caller already wrote directly into _hostBuffer via + // arrayBuffer, skip the redundant copy, the data is already in place. if (!(data instanceof ArrayBuffer && data === this._hostBuffer)) { this._writeToTarget(this._hostBuffer, data, options); } diff --git a/packages/typegpu/src/data/soaIO.ts b/packages/typegpu/src/data/soaIO.ts index 738934c846..60be48f577 100644 --- a/packages/typegpu/src/data/soaIO.ts +++ b/packages/typegpu/src/data/soaIO.ts @@ -69,6 +69,25 @@ export function isSoACompatibleField(schema: BaseData): boolean { return !isWgslStruct(schema); } +/** + * Computes the byte length that a scatterSoA call will naturally cover, based on + * the minimum element count implied by the provided SoA data arrays. + */ +export function computeSoAByteLength( + arraySchema: WgslArray, + soaData: Record, +): number | undefined { + const elementCount = inferSoAElementCount(arraySchema, soaData); + if (elementCount === undefined) { + return undefined; + } + const elementStride = roundUp( + sizeOf(arraySchema.elementType), + alignmentOf(arraySchema.elementType), + ); + return elementCount * elementStride; +} + export function getSoANaturalSize(dataType: BaseData, data: unknown): number | undefined { if ( !isWgslArray(dataType) || @@ -145,13 +164,9 @@ function writePackedValue( } /** - * Writes struct-of-arrays (SoA) data into a GPU-layout (AoS) target buffer. - * - * Each key in `soaData` is a struct field name mapped to a packed TypedArray - * containing that field's values for all elements (no inter-element padding). - * This function scatters those packed arrays into the correctly padded AoS layout. + * Scatters struct-of-arrays (SoA) data into a GPU-layout (AoS) target buffer. */ -export function writeSoA( +export function scatterSoA( target: Uint8Array, arraySchema: WgslArray, soaData: Record, @@ -161,6 +176,10 @@ export function writeSoA( const structSchema = arraySchema.elementType as WgslStruct; const offsets = offsetsForProps(structSchema); const elementStride = roundUp(sizeOf(structSchema), alignmentOf(structSchema)); + invariant( + startOffset % elementStride === 0, + `startOffset (${startOffset}) must be aligned to the element stride (${elementStride})`, + ); const startElement = Math.floor(startOffset / elementStride); const endElement = Math.min(arraySchema.elementCount, Math.ceil(endOffset / elementStride)); const elementCount = Math.max(0, endElement - startElement); @@ -171,9 +190,7 @@ export function writeSoA( continue; } const srcArray = soaData[key]; - if (srcArray === undefined) { - continue; - } + invariant(srcArray !== undefined, `Missing SoA data for field '${key}'`); const fieldOffset = offsets[key]?.offset; invariant(fieldOffset !== undefined, `Field ${key} not found in struct schema`); From 924417e57d8d645cd491e870aee6f4d47157d897 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 15:26:21 +0100 Subject: [PATCH 08/14] kill dead code --- packages/typegpu/src/data/soaIO.ts | 45 ++---------------------------- 1 file changed, 3 insertions(+), 42 deletions(-) diff --git a/packages/typegpu/src/data/soaIO.ts b/packages/typegpu/src/data/soaIO.ts index 60be48f577..0a0acaec21 100644 --- a/packages/typegpu/src/data/soaIO.ts +++ b/packages/typegpu/src/data/soaIO.ts @@ -4,7 +4,7 @@ import { alignmentOf } from './alignmentOf.ts'; import { offsetsForProps } from './offsets.ts'; import { sizeOf } from './sizeOf.ts'; import type { BaseData, WgslArray, WgslStruct } from './wgslTypes.ts'; -import { isMat, isMat2x2f, isMat3x3f, isWgslArray, isWgslStruct } from './wgslTypes.ts'; +import { isMat, isMat2x2f, isMat3x3f, isWgslArray } from './wgslTypes.ts'; function getPackedMatrixLayout(schema: BaseData) { if (!isMat(schema)) { @@ -21,7 +21,7 @@ function getPackedMatrixLayout(schema: BaseData) { } as const; } -export function packedSizeOf(schema: BaseData): number { +function packedSizeOf(schema: BaseData): number { const matrixLayout = getPackedMatrixLayout(schema); if (matrixLayout) { return matrixLayout.packedSize; @@ -34,7 +34,7 @@ export function packedSizeOf(schema: BaseData): number { return sizeOf(schema); } -export function inferSoAElementCount( +function inferSoAElementCount( arraySchema: WgslArray, soaData: Record, ): number | undefined { @@ -61,14 +61,6 @@ export function inferSoAElementCount( return inferredCount; } -export function isSoACompatibleField(schema: BaseData): boolean { - if (isWgslArray(schema)) { - return isSoACompatibleField(schema.elementType); - } - - return !isWgslStruct(schema); -} - /** * Computes the byte length that a scatterSoA call will naturally cover, based on * the minimum element count implied by the provided SoA data arrays. @@ -88,37 +80,6 @@ export function computeSoAByteLength( return elementCount * elementStride; } -export function getSoANaturalSize(dataType: BaseData, data: unknown): number | undefined { - if ( - !isWgslArray(dataType) || - !isWgslStruct(dataType.elementType) || - Array.isArray(data) || - typeof data !== 'object' || - data === null - ) { - return undefined; - } - - const soaData = data as Record; - const values = Object.values(soaData); - const isSoAInput = - values.length > 0 && - values.every((value) => ArrayBuffer.isView(value)) && - Object.values(dataType.elementType.propTypes).every(isSoACompatibleField); - - if (!isSoAInput) { - return undefined; - } - - const elementCount = inferSoAElementCount(dataType, soaData as Record); - if (elementCount === undefined) { - return undefined; - } - - const elementStride = roundUp(sizeOf(dataType.elementType), alignmentOf(dataType.elementType)); - return elementCount * elementStride; -} - function writePackedValue( target: Uint8Array, schema: BaseData, From a0aa1e4eaf1cf8eedcf3fe04754fefc9dca1dce7 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 18:20:42 +0100 Subject: [PATCH 09/14] fixes --- .../src/content/docs/fundamentals/buffers.mdx | 2 +- packages/typegpu/src/common/writeSoA.ts | 163 ++++++++++++++++- packages/typegpu/src/core/buffer/buffer.ts | 30 +-- packages/typegpu/src/data/soaIO.ts | 171 ------------------ 4 files changed, 178 insertions(+), 188 deletions(-) delete mode 100644 packages/typegpu/src/data/soaIO.ts diff --git a/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx b/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx index 97950077c4..4352c1d9f5 100644 --- a/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx +++ b/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx @@ -29,7 +29,7 @@ const Particle = d.struct({ }); // Utility for creating a random particle -function createParticle(): d.Infer { +function createParticle(): d.InferInput { return { position: d.vec3f(Math.random(), 2, Math.random()), velocity: d.vec3f(0, 9.8, 0), diff --git a/packages/typegpu/src/common/writeSoA.ts b/packages/typegpu/src/common/writeSoA.ts index 368af6ab5f..047e782dbf 100644 --- a/packages/typegpu/src/common/writeSoA.ts +++ b/packages/typegpu/src/common/writeSoA.ts @@ -1,8 +1,169 @@ -import { computeSoAByteLength, scatterSoA } from '../data/soaIO.ts'; +import { invariant } from '../errors.ts'; +import { roundUp } from '../mathUtils.ts'; +import { alignmentOf } from '../data/alignmentOf.ts'; +import { offsetsForProps } from '../data/offsets.ts'; import { sizeOf } from '../data/sizeOf.ts'; import type { BaseData, SoAInputFor, WgslArray, WgslStruct } from '../data/wgslTypes.ts'; +import { isMat, isMat2x2f, isMat3x3f, isWgslArray } from '../data/wgslTypes.ts'; import type { BufferWriteOptions, TgpuBuffer } from '../core/buffer/buffer.ts'; +function getPackedMatrixLayout(schema: BaseData) { + if (!isMat(schema)) { + return undefined; + } + + const dim = isMat3x3f(schema) ? 3 : isMat2x2f(schema) ? 2 : 4; + const packedColumnSize = dim * 4; + + return { + dim, + packedColumnSize, + packedSize: dim * packedColumnSize, + } as const; +} + +function packedSizeOf(schema: BaseData): number { + const matrixLayout = getPackedMatrixLayout(schema); + if (matrixLayout) { + return matrixLayout.packedSize; + } + + if (isWgslArray(schema)) { + return schema.elementCount * packedSizeOf(schema.elementType); + } + + return sizeOf(schema); +} + +function inferSoAElementCount( + arraySchema: WgslArray, + soaData: Record, +): number | undefined { + const structSchema = arraySchema.elementType as WgslStruct; + let inferredCount: number | undefined; + + for (const key in soaData) { + const srcArray = soaData[key]; + const fieldSchema = structSchema.propTypes[key]; + if (srcArray === undefined || fieldSchema === undefined) { + continue; + } + + const fieldPackedSize = packedSizeOf(fieldSchema); + if (fieldPackedSize === 0) { + continue; + } + + const fieldElementCount = Math.floor(srcArray.byteLength / fieldPackedSize); + inferredCount = + inferredCount === undefined ? fieldElementCount : Math.min(inferredCount, fieldElementCount); + } + + return inferredCount; +} + +function computeSoAByteLength( + arraySchema: WgslArray, + soaData: Record, +): number | undefined { + const elementCount = inferSoAElementCount(arraySchema, soaData); + if (elementCount === undefined) { + return undefined; + } + const elementStride = roundUp( + sizeOf(arraySchema.elementType), + alignmentOf(arraySchema.elementType), + ); + return elementCount * elementStride; +} + +function writePackedValue( + target: Uint8Array, + schema: BaseData, + srcBytes: Uint8Array, + dstOffset: number, + srcOffset: number, +): void { + const matrixLayout = getPackedMatrixLayout(schema); + if (matrixLayout) { + const gpuColumnStride = roundUp(matrixLayout.packedColumnSize, alignmentOf(schema)); + + for (let col = 0; col < matrixLayout.dim; col++) { + target.set( + srcBytes.subarray( + srcOffset + col * matrixLayout.packedColumnSize, + srcOffset + col * matrixLayout.packedColumnSize + matrixLayout.packedColumnSize, + ), + dstOffset + col * gpuColumnStride, + ); + } + + return; + } + + if (isWgslArray(schema)) { + const packedElementSize = packedSizeOf(schema.elementType); + const gpuElementStride = roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType)); + + for (let i = 0; i < schema.elementCount; i++) { + writePackedValue( + target, + schema.elementType, + srcBytes, + dstOffset + i * gpuElementStride, + srcOffset + i * packedElementSize, + ); + } + + return; + } + + target.set(srcBytes.subarray(srcOffset, srcOffset + sizeOf(schema)), dstOffset); +} + +function scatterSoA( + target: Uint8Array, + arraySchema: WgslArray, + soaData: Record, + startOffset: number, + endOffset: number, +): void { + const structSchema = arraySchema.elementType as WgslStruct; + const offsets = offsetsForProps(structSchema); + const elementStride = roundUp(sizeOf(structSchema), alignmentOf(structSchema)); + invariant( + startOffset % elementStride === 0, + `startOffset (${startOffset}) must be aligned to the element stride (${elementStride})`, + ); + const startElement = Math.floor(startOffset / elementStride); + const endElement = Math.min(arraySchema.elementCount, Math.ceil(endOffset / elementStride)); + const elementCount = Math.max(0, endElement - startElement); + + for (const key in structSchema.propTypes) { + const fieldSchema = structSchema.propTypes[key]; + if (fieldSchema === undefined) { + continue; + } + const srcArray = soaData[key]; + invariant(srcArray !== undefined, `Missing SoA data for field '${key}'`); + + const fieldOffset = offsets[key]?.offset; + invariant(fieldOffset !== undefined, `Field ${key} not found in struct schema`); + const srcBytes = new Uint8Array(srcArray.buffer, srcArray.byteOffset, srcArray.byteLength); + + const packedFieldSize = packedSizeOf(fieldSchema); + for (let i = 0; i < elementCount; i++) { + writePackedValue( + target, + fieldSchema, + srcBytes, + (startElement + i) * elementStride + fieldOffset, + i * packedFieldSize, + ); + } + } +} + export function writeSoA>( buffer: TgpuBuffer>>, data: SoAInputFor, diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index bed001f91b..8982838983 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -145,7 +145,7 @@ export interface TgpuBuffer extends TgpuNamable { as>(usage: T): UsageTypeToBufferUsage[T]; compileWriter(): void; - write(data: Prettify>, options?: BufferWriteOptions): void; + write(data: InferInput, options?: BufferWriteOptions): void; write(data: ArrayBuffer, options?: BufferWriteOptions): void; writePartial(data: InferPartial): void; clear(): void; @@ -246,9 +246,9 @@ class TgpuBufferImpl implements TgpuBuffer { if (this._initialCallback) { this._initialCallback(this); } else if (this.initial) { - this._writeToTarget(this._getMappedRange(), this.initial); + this.#writeToTarget(this.#getMappedRange(), this.initial); } - this._unmapBuffer(); + this.#unmapBuffer(); } } @@ -262,7 +262,7 @@ class TgpuBufferImpl implements TgpuBuffer { get arrayBuffer(): ArrayBuffer { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - return this._getMappedRange(); + return this.#getMappedRange(); } if (!this._hostBuffer) { @@ -272,7 +272,7 @@ class TgpuBufferImpl implements TgpuBuffer { return this._hostBuffer; } - private _getMappedRange(): ArrayBuffer { + #getMappedRange(): ArrayBuffer { if (!this._buffer || this._buffer.mapState !== 'mapped') { throw new Error('Buffer is not mapped.'); } @@ -281,7 +281,7 @@ class TgpuBufferImpl implements TgpuBuffer { return this._mappedRange; } - private _unmapBuffer(): void { + #unmapBuffer(): void { if (!this._buffer || this._buffer.mapState !== 'mapped') { return; } @@ -343,7 +343,7 @@ class TgpuBufferImpl implements TgpuBuffer { getCompiledWriterForSchema(this.dataType); } - private _writeToTarget( + #writeToTarget( target: ArrayBuffer, data: InferInput | ArrayBuffer, options?: BufferWriteOptions, @@ -416,13 +416,13 @@ class TgpuBufferImpl implements TgpuBuffer { const size = endOffset - startOffset; if (gpuBuffer.mapState === 'mapped') { - const mapped = this._getMappedRange(); + const mapped = this.#getMappedRange(); if (data instanceof ArrayBuffer && data === mapped) { // The caller already wrote data directly into the mapped range // via arrayBuffer. Nothing to do here return; } - this._writeToTarget(mapped, data, options); + this.#writeToTarget(mapped, data, options); return; } @@ -433,7 +433,7 @@ class TgpuBufferImpl implements TgpuBuffer { // If the caller already wrote directly into _hostBuffer via // arrayBuffer, skip the redundant copy, the data is already in place. if (!(data instanceof ArrayBuffer && data === this._hostBuffer)) { - this._writeToTarget(this._hostBuffer, data, options); + this.#writeToTarget(this._hostBuffer, data, options); } this.#device.queue.writeBuffer(gpuBuffer, startOffset, this._hostBuffer, startOffset, size); } @@ -444,7 +444,7 @@ class TgpuBufferImpl implements TgpuBuffer { const instructions = getWriteInstructions(this.dataType, data); if (gpuBuffer.mapState === 'mapped') { - const mappedRange = this._getMappedRange(); + const mappedRange = this.#getMappedRange(); const mappedView = new Uint8Array(mappedRange); for (const instruction of instructions) { @@ -467,7 +467,7 @@ class TgpuBufferImpl implements TgpuBuffer { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - new Uint8Array(this._getMappedRange()).fill(0); + new Uint8Array(this.#getMappedRange()).fill(0); return; } @@ -491,15 +491,15 @@ class TgpuBufferImpl implements TgpuBuffer { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - const mapped = this._getMappedRange(); + const mapped = this.#getMappedRange(); return readData(new BufferReader(mapped), this.dataType); } if (gpuBuffer.usage & GPUBufferUsage.MAP_READ) { await gpuBuffer.mapAsync(GPUMapMode.READ); - const mapped = this._getMappedRange(); + const mapped = this.#getMappedRange(); const res = readData(new BufferReader(mapped), this.dataType); - this._unmapBuffer(); + this.#unmapBuffer(); return res; } diff --git a/packages/typegpu/src/data/soaIO.ts b/packages/typegpu/src/data/soaIO.ts deleted file mode 100644 index 0a0acaec21..0000000000 --- a/packages/typegpu/src/data/soaIO.ts +++ /dev/null @@ -1,171 +0,0 @@ -import { invariant } from '../errors.ts'; -import { roundUp } from '../mathUtils.ts'; -import { alignmentOf } from './alignmentOf.ts'; -import { offsetsForProps } from './offsets.ts'; -import { sizeOf } from './sizeOf.ts'; -import type { BaseData, WgslArray, WgslStruct } from './wgslTypes.ts'; -import { isMat, isMat2x2f, isMat3x3f, isWgslArray } from './wgslTypes.ts'; - -function getPackedMatrixLayout(schema: BaseData) { - if (!isMat(schema)) { - return undefined; - } - - const dim = isMat3x3f(schema) ? 3 : isMat2x2f(schema) ? 2 : 4; - const packedColumnSize = dim * 4; - - return { - dim, - packedColumnSize, - packedSize: dim * packedColumnSize, - } as const; -} - -function packedSizeOf(schema: BaseData): number { - const matrixLayout = getPackedMatrixLayout(schema); - if (matrixLayout) { - return matrixLayout.packedSize; - } - - if (isWgslArray(schema)) { - return schema.elementCount * packedSizeOf(schema.elementType); - } - - return sizeOf(schema); -} - -function inferSoAElementCount( - arraySchema: WgslArray, - soaData: Record, -): number | undefined { - const structSchema = arraySchema.elementType as WgslStruct; - let inferredCount: number | undefined; - - for (const key in soaData) { - const srcArray = soaData[key]; - const fieldSchema = structSchema.propTypes[key]; - if (srcArray === undefined || fieldSchema === undefined) { - continue; - } - - const fieldPackedSize = packedSizeOf(fieldSchema); - if (fieldPackedSize === 0) { - continue; - } - - const fieldElementCount = Math.floor(srcArray.byteLength / fieldPackedSize); - inferredCount = - inferredCount === undefined ? fieldElementCount : Math.min(inferredCount, fieldElementCount); - } - - return inferredCount; -} - -/** - * Computes the byte length that a scatterSoA call will naturally cover, based on - * the minimum element count implied by the provided SoA data arrays. - */ -export function computeSoAByteLength( - arraySchema: WgslArray, - soaData: Record, -): number | undefined { - const elementCount = inferSoAElementCount(arraySchema, soaData); - if (elementCount === undefined) { - return undefined; - } - const elementStride = roundUp( - sizeOf(arraySchema.elementType), - alignmentOf(arraySchema.elementType), - ); - return elementCount * elementStride; -} - -function writePackedValue( - target: Uint8Array, - schema: BaseData, - srcBytes: Uint8Array, - dstOffset: number, - srcOffset: number, -): void { - const matrixLayout = getPackedMatrixLayout(schema); - if (matrixLayout) { - const gpuColumnStride = roundUp(matrixLayout.packedColumnSize, alignmentOf(schema)); - - for (let col = 0; col < matrixLayout.dim; col++) { - target.set( - srcBytes.subarray( - srcOffset + col * matrixLayout.packedColumnSize, - srcOffset + col * matrixLayout.packedColumnSize + matrixLayout.packedColumnSize, - ), - dstOffset + col * gpuColumnStride, - ); - } - - return; - } - - if (isWgslArray(schema)) { - const packedElementSize = packedSizeOf(schema.elementType); - const gpuElementStride = roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType)); - - for (let i = 0; i < schema.elementCount; i++) { - writePackedValue( - target, - schema.elementType, - srcBytes, - dstOffset + i * gpuElementStride, - srcOffset + i * packedElementSize, - ); - } - - return; - } - - target.set(srcBytes.subarray(srcOffset, srcOffset + sizeOf(schema)), dstOffset); -} - -/** - * Scatters struct-of-arrays (SoA) data into a GPU-layout (AoS) target buffer. - */ -export function scatterSoA( - target: Uint8Array, - arraySchema: WgslArray, - soaData: Record, - startOffset: number, - endOffset: number, -): void { - const structSchema = arraySchema.elementType as WgslStruct; - const offsets = offsetsForProps(structSchema); - const elementStride = roundUp(sizeOf(structSchema), alignmentOf(structSchema)); - invariant( - startOffset % elementStride === 0, - `startOffset (${startOffset}) must be aligned to the element stride (${elementStride})`, - ); - const startElement = Math.floor(startOffset / elementStride); - const endElement = Math.min(arraySchema.elementCount, Math.ceil(endOffset / elementStride)); - const elementCount = Math.max(0, endElement - startElement); - - for (const key in structSchema.propTypes) { - const fieldSchema = structSchema.propTypes[key]; - if (fieldSchema === undefined) { - continue; - } - const srcArray = soaData[key]; - invariant(srcArray !== undefined, `Missing SoA data for field '${key}'`); - - const fieldOffset = offsets[key]?.offset; - invariant(fieldOffset !== undefined, `Field ${key} not found in struct schema`); - const srcBytes = new Uint8Array(srcArray.buffer, srcArray.byteOffset, srcArray.byteLength); - - const packedFieldSize = packedSizeOf(fieldSchema); - for (let i = 0; i < elementCount; i++) { - writePackedValue( - target, - fieldSchema, - srcBytes, - (startElement + i) * elementStride + fieldOffset, - i * packedFieldSize, - ); - } - } -} From 03c8303cde6e030edec45b96214774f554df07e6 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 18:45:35 +0100 Subject: [PATCH 10/14] docs --- .../src/content/docs/fundamentals/buffers.mdx | 107 ++++++++++++++++++ packages/typegpu/tests/buffer.test.ts | 37 ++++++ 2 files changed, 144 insertions(+) diff --git a/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx b/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx index 4352c1d9f5..b8e019d5ca 100644 --- a/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx +++ b/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx @@ -160,6 +160,26 @@ const buffer2 = root.createBuffer(d.arrayOf(d.vec3f, 2), [ ]); ``` +For cases where a plain typed value is not enough, you can also pass an initializer callback. +It receives the newly created typed buffer while it is still mapped, so you can populate it with multiple writes or helper utilities before the first upload. + +```ts twoslash +import tgpu, { d } from 'typegpu'; + +const root = await tgpu.init(); +// ---cut--- +const Schema = d.arrayOf(d.u32, 6); +const firstChunk = d.memoryLayoutOf(Schema, (a) => a[1]); +const secondChunk = d.memoryLayoutOf(Schema, (a) => a[4]); + +const buffer = root.createBuffer(Schema, (mappedBuffer) => { + mappedBuffer.write([10, 20], { startOffset: firstChunk.offset }); + mappedBuffer.write([30, 40], { startOffset: secondChunk.offset }); +}); +``` + +For buffers whose CPU-side data already lives in separate per-field arrays, see the SoA write section. + ### Using an existing buffer You can also create a buffer using an existing WebGPU buffer. This is useful when you have existing logic but want to introduce type-safe data operations. @@ -383,6 +403,93 @@ planetBuffer.writePartial({ }); ``` +### Writing struct-of-arrays (SoA) data + +When the buffer schema is an `array>`, you can write the data in a struct-of-arrays form with `writeSoA` from `typegpu/common`. +This is useful when your CPU-side data is already stored per-field, such as simulation attributes kept in separate typed arrays. + +```ts twoslash +import tgpu, { d } from 'typegpu'; +import { writeSoA } from 'typegpu/common'; + +const root = await tgpu.init(); +// ---cut--- +const Particle = d.struct({ + pos: d.vec3f, + vel: d.f32, +}); + +const particleBuffer = root.createBuffer(d.arrayOf(Particle, 2)); + +writeSoA(particleBuffer, { + pos: new Float32Array([ + 1, 2, 3, + 4, 5, 6, + ]), + vel: new Float32Array([10, 20]), +}); +``` + +`writeSoA` converts those field-wise arrays into the buffer's GPU array-of-structs layout before uploading. +For the example above, the resulting bytes match two `Particle` structs laid out in memory, including the padding required by `vec3f`. + +:::note[Supported field shapes] +SoA writes work for arrays of structs whose fields are: +- scalars, +- vectors, +- matrices, +- fixed-size arrays of scalars, vectors, or matrices. + +Nested struct fields are not supported. +::: + +:::tip[Packed input, padded output] +Unlike `.write()` with a raw `TypedArray` or `ArrayBuffer`, `writeSoA` expects each field in its natural packed form. +For example, a `vec3f` field should be provided as 3 floats per element, not 4, and a `mat3x3f` field should be provided as 9 floats per matrix, not 12. +TypeGPU inserts the required WGSL padding while scattering the data into the target buffer. +::: + +You can also restrict the write to a slice of the destination buffer: + +```ts twoslash +import tgpu, { d } from 'typegpu'; +import { writeSoA } from 'typegpu/common'; + +const root = await tgpu.init(); +// ---cut--- +const Entry = d.struct({ + id: d.u32, + values: d.arrayOf(d.vec3f, 2), +}); + +const schema = d.arrayOf(Entry, 4); +const buffer = root.createBuffer(schema); + +const start = d.memoryLayoutOf(schema, (a) => a[1]); +const end = d.memoryLayoutOf(schema, (a) => a[3]); + +writeSoA( + buffer, + { + id: new Uint32Array([30, 40]), + values: new Float32Array([ + 1, 2, 3, + 4, 5, 6, + 7, 8, 9, + 10, 11, 12, + ]), + }, + { + startOffset: start.offset, + endOffset: end.offset, + }, +); +``` + +If `endOffset` is omitted, TypeGPU infers the written range from the provided field arrays and stops at the shortest complete element count implied by the data. +`startOffset` is still byte-based and should point to the start of a struct element, so `d.memoryLayoutOf` is the safest way to compute it. +All struct fields must be present in the SoA object. + ### Copying There's also an option to copy value from another typed buffer using the `.copyFrom(buffer)` method, diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index a5fff10ba0..8bb29a670d 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -80,6 +80,16 @@ describe('TgpuBuffer', () => { const s1 = d.struct({ a: d.u32, b: d.u32, c: d.vec3i }); const s2 = d.struct({ a: d.u32, b: s1, c: d.vec4u }); + const schema = d.arrayOf(d.u32, 6); + const firstChunk = d.memoryLayoutOf(schema, (a) => a[1]); + const secondChunk = d.memoryLayoutOf(schema, (a) => a[4]); + + const buffer = root.createBuffer(schema, (mappedBuffer) => { + mappedBuffer.write([10, 20], { startOffset: firstChunk.offset }); + mappedBuffer.write([30, 40], { startOffset: secondChunk.offset }); + }); + const raw = buffer.buffer; + const dataBuffer = root.createBuffer(s2).$usage('uniform'); root.unwrap(dataBuffer); @@ -1163,6 +1173,33 @@ describe('ValidateBufferSchema', () => { expect([...result]).toStrictEqual([1, 2, 3, 0, 4, 5, 6, 0, 7, 8, 9, 0, 10, 11, 12, 0]); }); + it('should write SoA data for struct fields that are arrays of arrays of padded vectors', ({ + root, + device, + }) => { + const Entry = d.struct({ + values: d.arrayOf(d.arrayOf(d.vec3f, 2), 2), + }); + + const schema = d.arrayOf(Entry, 2); + const buffer = root.createBuffer(schema); + root.unwrap(buffer); + + common.writeSoA(buffer, { + values: new Float32Array([ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, + ]), + }); + + const uploadedBuffer = device.mock.queue.writeBuffer.mock.calls[0]?.[2] as ArrayBuffer; + const result = new Float32Array(uploadedBuffer); + + expect([...result]).toStrictEqual([ + 1, 2, 3, 0, 4, 5, 6, 0, 7, 8, 9, 0, 10, 11, 12, 0, 13, 14, 15, 0, 16, 17, 18, 0, 19, 20, 21, + 0, 22, 23, 24, 0, + ]); + }); + it('should write SoA data for struct fields that are arrays of padded matrices', ({ root, device, From 76a677f5ac22b4666badee166f4293f30acbeb94 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 18:46:26 +0100 Subject: [PATCH 11/14] remove debug --- packages/typegpu/tests/buffer.test.ts | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index 8bb29a670d..7d0debae4e 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -80,16 +80,6 @@ describe('TgpuBuffer', () => { const s1 = d.struct({ a: d.u32, b: d.u32, c: d.vec3i }); const s2 = d.struct({ a: d.u32, b: s1, c: d.vec4u }); - const schema = d.arrayOf(d.u32, 6); - const firstChunk = d.memoryLayoutOf(schema, (a) => a[1]); - const secondChunk = d.memoryLayoutOf(schema, (a) => a[4]); - - const buffer = root.createBuffer(schema, (mappedBuffer) => { - mappedBuffer.write([10, 20], { startOffset: firstChunk.offset }); - mappedBuffer.write([30, 40], { startOffset: secondChunk.offset }); - }); - const raw = buffer.buffer; - const dataBuffer = root.createBuffer(s2).$usage('uniform'); root.unwrap(dataBuffer); From 0acfe4d196894a503c57983fca2fa90867e0af32 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 18:48:41 +0100 Subject: [PATCH 12/14] better link --- apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx b/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx index b8e019d5ca..cf989da042 100644 --- a/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx +++ b/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx @@ -178,7 +178,7 @@ const buffer = root.createBuffer(Schema, (mappedBuffer) => { }); ``` -For buffers whose CPU-side data already lives in separate per-field arrays, see the SoA write section. +For buffers whose CPU-side data already lives in separate per-field arrays, see the [SoA write section](#writing-struct-of-arrays-soa-data). ### Using an existing buffer From 20d753854b5094d97ed4daf5c4b7a3039b61b307 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 18:50:35 +0100 Subject: [PATCH 13/14] better import --- .../src/content/docs/fundamentals/buffers.mdx | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx b/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx index cf989da042..73e95ca1af 100644 --- a/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx +++ b/apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx @@ -409,8 +409,7 @@ When the buffer schema is an `array>`, you can write the data in a s This is useful when your CPU-side data is already stored per-field, such as simulation attributes kept in separate typed arrays. ```ts twoslash -import tgpu, { d } from 'typegpu'; -import { writeSoA } from 'typegpu/common'; +import tgpu, { d, common } from 'typegpu'; const root = await tgpu.init(); // ---cut--- @@ -421,7 +420,7 @@ const Particle = d.struct({ const particleBuffer = root.createBuffer(d.arrayOf(Particle, 2)); -writeSoA(particleBuffer, { +common.writeSoA(particleBuffer, { pos: new Float32Array([ 1, 2, 3, 4, 5, 6, @@ -452,8 +451,7 @@ TypeGPU inserts the required WGSL padding while scattering the data into the tar You can also restrict the write to a slice of the destination buffer: ```ts twoslash -import tgpu, { d } from 'typegpu'; -import { writeSoA } from 'typegpu/common'; +import tgpu, { d, common } from 'typegpu'; const root = await tgpu.init(); // ---cut--- @@ -468,7 +466,7 @@ const buffer = root.createBuffer(schema); const start = d.memoryLayoutOf(schema, (a) => a[1]); const end = d.memoryLayoutOf(schema, (a) => a[3]); -writeSoA( +common.writeSoA( buffer, { id: new Uint32Array([30, 40]), From a151260cd4dbedb0cf99b0327f786e9a455967d7 Mon Sep 17 00:00:00 2001 From: Konrad Reczko Date: Fri, 27 Mar 2026 19:04:02 +0100 Subject: [PATCH 14/14] cleaner tyep --- packages/typegpu/src/common/writeSoA.ts | 18 +++++++++++++++++- packages/typegpu/src/data/index.ts | 1 - packages/typegpu/src/data/wgslTypes.ts | 16 ---------------- packages/typegpu/tests/buffer.test.ts | 5 ++--- 4 files changed, 19 insertions(+), 21 deletions(-) diff --git a/packages/typegpu/src/common/writeSoA.ts b/packages/typegpu/src/common/writeSoA.ts index 047e782dbf..a867373b92 100644 --- a/packages/typegpu/src/common/writeSoA.ts +++ b/packages/typegpu/src/common/writeSoA.ts @@ -3,9 +3,21 @@ import { roundUp } from '../mathUtils.ts'; import { alignmentOf } from '../data/alignmentOf.ts'; import { offsetsForProps } from '../data/offsets.ts'; import { sizeOf } from '../data/sizeOf.ts'; -import type { BaseData, SoAInputFor, WgslArray, WgslStruct } from '../data/wgslTypes.ts'; +import type { BaseData, TypedArrayFor, WgslArray, WgslStruct } from '../data/wgslTypes.ts'; import { isMat, isMat2x2f, isMat3x3f, isWgslArray } from '../data/wgslTypes.ts'; import type { BufferWriteOptions, TgpuBuffer } from '../core/buffer/buffer.ts'; +import type { Prettify } from '../shared/utilityTypes.ts'; + +type UnwrapWgslArray = T extends WgslArray ? UnwrapWgslArray : T; +type PackedSoAInputFor = TypedArrayFor>; + +type SoAFieldsFor> = { + [K in keyof T as [PackedSoAInputFor] extends [never] ? never : K]: PackedSoAInputFor; +}; + +type SoAInputFor> = [keyof T] extends [keyof SoAFieldsFor] + ? Prettify> + : never; function getPackedMatrixLayout(schema: BaseData) { if (!isMat(schema)) { @@ -189,3 +201,7 @@ export function writeSoA>( ); buffer.write(arrayBuffer, { startOffset, endOffset }); } + +export namespace writeSoA { + export type InputFor> = SoAInputFor; +} diff --git a/packages/typegpu/src/data/index.ts b/packages/typegpu/src/data/index.ts index 7cd58ca2ea..7b51b0467c 100644 --- a/packages/typegpu/src/data/index.ts +++ b/packages/typegpu/src/data/index.ts @@ -89,7 +89,6 @@ export type { Mat4x4f, matBase, Ptr, - SoAInputFor, Size, StorableData, U16, diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts index a6738e7452..9a71fc252b 100644 --- a/packages/typegpu/src/data/wgslTypes.ts +++ b/packages/typegpu/src/data/wgslTypes.ts @@ -67,22 +67,6 @@ export type TypedArrayFor = T extends F32 | Vec2f | Vec3f | Vec4f | Mat2x2f | ? Uint16Array : never; -type UnwrapWgslArray = T extends WgslArray ? UnwrapWgslArray : T; -type PackedSoAInputFor = TypedArrayFor>; - -type SoAFieldsFor> = { - [K in keyof T as [PackedSoAInputFor] extends [never] ? never : K]: PackedSoAInputFor; -}; - -/** - * Maps struct properties to a record of TypedArrays (Struct-of-Arrays input format) - if not possible, resolves to never. - */ -export type SoAInputFor> = [keyof T] extends [ - keyof SoAFieldsFor, -] - ? Prettify> - : never; - /** * Vector infix notation. * diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index 7d0debae4e..14cff19735 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -2,7 +2,6 @@ import { attest } from '@ark/attest'; import { describe, expect, expectTypeOf, vi } from 'vitest'; import * as common from '../src/common/index.ts'; import * as d from '../src/data/index.ts'; -import type { InferInput, SoAInputFor } from '../src/data/index.ts'; import { sizeOf } from '../src/data/sizeOf.ts'; import type { ValidateBufferSchema, ValidUsagesFor } from '../src/index.js'; import { getName } from '../src/shared/meta.ts'; @@ -1085,7 +1084,7 @@ describe('ValidateBufferSchema', () => { f: d.WgslArray; }; - expectTypeOf>().toEqualTypeOf<{ + expectTypeOf>().toEqualTypeOf<{ a: Float32Array; b: Uint32Array; c: Float32Array; @@ -1105,7 +1104,7 @@ describe('ValidateBufferSchema', () => { nested: typeof Nested; }; - expectTypeOf>().toEqualTypeOf(); + expectTypeOf>().toEqualTypeOf(); }); it('should write SoA data for struct fields that are fixed-size arrays of primitives', ({