From c0c4088e42ac62fa51ef5aa8f48af2488f83af46 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Mon, 19 Oct 2026 12:42:54 +0200 Subject: [PATCH 1/4] strideOf helper --- packages/typegpu/src/common/writeSoA.ts | 10 +++--- .../src/core/pipeline/webgpuLimitations.ts | 3 +- .../src/core/vertexLayout/vertexLayout.ts | 3 +- packages/typegpu/src/data/compiledIO.ts | 4 +-- packages/typegpu/src/data/dataIO.ts | 5 ++- packages/typegpu/src/data/offsetUtils.ts | 5 ++- packages/typegpu/src/data/partialIO.ts | 5 ++- .../typegpu/src/data/schemaMemoryLayout.ts | 31 +++++++++++++------ 8 files changed, 37 insertions(+), 29 deletions(-) diff --git a/packages/typegpu/src/common/writeSoA.ts b/packages/typegpu/src/common/writeSoA.ts index 783e1107aa..47e195fbb0 100644 --- a/packages/typegpu/src/common/writeSoA.ts +++ b/packages/typegpu/src/common/writeSoA.ts @@ -5,6 +5,7 @@ import { alignmentOf } from '../data/alignmentOf.ts'; import { undecorate } from '../data/dataTypes.ts'; import { offsetsForProps } from '../data/offsets.ts'; import { sizeOf } from '../data/sizeOf.ts'; +import { strideOf } from '../data/schemaMemoryLayout.ts'; import type { BaseData, TypedArrayFor, WgslArray, WgslStruct } from '../data/wgslTypes.ts'; import { isAtomic, isMat, isMat2x2f, isMat3x3f, isWgslArray } from '../data/wgslTypes.ts'; import type { BufferWriteOptions, TgpuBuffer } from '../core/buffer/buffer.ts'; @@ -69,7 +70,7 @@ function computeSoAByteLength( if (inferredCount === undefined) { return undefined; } - const elementStride = roundUp(sizeOf(structSchema), alignmentOf(structSchema)); + const elementStride = strideOf(arraySchema); return inferredCount * elementStride; } @@ -99,10 +100,7 @@ function writePackedValue( } if (isWgslArray(unpackedSchema)) { const packedElementSize = packedSizeOf(unpackedSchema.elementType); - const gpuElementStride = roundUp( - sizeOf(unpackedSchema.elementType), - alignmentOf(unpackedSchema.elementType), - ); + const gpuElementStride = strideOf(unpackedSchema); for (let i = 0; i < unpackedSchema.elementCount; i++) { writePackedValue( @@ -126,7 +124,7 @@ function scatterSoA( endOffset: number, ): void { const structSchema = arraySchema.elementType as WgslStruct; - const elementStride = roundUp(sizeOf(structSchema), alignmentOf(structSchema)); + const elementStride = strideOf(arraySchema); invariant( startOffset % elementStride === 0, `startOffset (${startOffset}) must be aligned to the element stride (${elementStride})`, diff --git a/packages/typegpu/src/core/pipeline/webgpuLimitations.ts b/packages/typegpu/src/core/pipeline/webgpuLimitations.ts index 64eb46e672..2d125870f4 100644 --- a/packages/typegpu/src/core/pipeline/webgpuLimitations.ts +++ b/packages/typegpu/src/core/pipeline/webgpuLimitations.ts @@ -1,6 +1,7 @@ import { alignmentOf } from '../../data/alignmentOf.ts'; import { memoryLayoutOf } from '../../data/offsetUtils.ts'; import { sizeOf } from '../../data/sizeOf.ts'; +import { strideOf } from '../../data/schemaMemoryLayout.ts'; import { isWgslArray, isWgslStruct, type BaseData } from '../../data/wgslTypes.ts'; import { invariant } from '../../errors.ts'; import { roundUp } from '../../mathUtils.ts'; @@ -59,7 +60,7 @@ function checkUniformAlignment(schema: BaseData) { if (isWgslArray(schema)) { checkUniformAlignment(schema.elementType); - const stride = roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType)); + const stride = strideOf(schema); if (stride % 16 !== 0) { logger.warn( 'uniform-schema-misaligned', diff --git a/packages/typegpu/src/core/vertexLayout/vertexLayout.ts b/packages/typegpu/src/core/vertexLayout/vertexLayout.ts index 769d142011..f2478e2a2c 100644 --- a/packages/typegpu/src/core/vertexLayout/vertexLayout.ts +++ b/packages/typegpu/src/core/vertexLayout/vertexLayout.ts @@ -2,6 +2,7 @@ import { alignmentOf, customAlignmentOf } from '../../data/alignmentOf.ts'; import type { Disarray } from '../../data/dataTypes.ts'; import { getCustomLocation, isLooseDecorated, isUnstruct } from '../../data/dataTypes.ts'; import { sizeOf } from '../../data/sizeOf.ts'; +import { strideOf } from '../../data/schemaMemoryLayout.ts'; import type { BaseData, WgslArray } from '../../data/wgslTypes.ts'; import { isDecorated, isWgslStruct } from '../../data/wgslTypes.ts'; import { roundUp } from '../../mathUtils.ts'; @@ -161,7 +162,7 @@ class TgpuVertexLayoutImpl implements TgpuVe stepMode, label: undefined, }; - this.stride = roundUp(sizeOf(arraySchema.elementType), alignmentOf(arraySchema)); + this.stride = strideOf(arraySchema); this.attrib = dataToContainedAttribs(this, arraySchema.elementType, 0, this.#customLocationMap); } diff --git a/packages/typegpu/src/data/compiledIO.ts b/packages/typegpu/src/data/compiledIO.ts index 7d510add40..354984d8b3 100644 --- a/packages/typegpu/src/data/compiledIO.ts +++ b/packages/typegpu/src/data/compiledIO.ts @@ -1,9 +1,9 @@ import { roundUp } from '../mathUtils.ts'; import { logger } from '../tgpuLogger.ts'; -import { alignmentOf } from './alignmentOf.ts'; import { isDisarray, isLooseDecorated, isUnstruct } from './dataTypes.ts'; import { offsetsForProps } from './offsets.ts'; import { sizeOf } from './sizeOf.ts'; +import { strideOf } from './schemaMemoryLayout.ts'; import { formatToWGSLType, isPackedData } from './vertexFormatData.ts'; import * as wgsl from './wgslTypes.ts'; @@ -186,7 +186,7 @@ export function buildWriter( throw new Error('Cannot write using a runtime-sized schema.'); } - const elementSize = roundUp(sizeOf(node.elementType), alignmentOf(node)); + const elementSize = strideOf(node); const totalSize = node.elementCount * elementSize; const copyLen = partial diff --git a/packages/typegpu/src/data/dataIO.ts b/packages/typegpu/src/data/dataIO.ts index 1d4a75cce7..4b70bbe724 100644 --- a/packages/typegpu/src/data/dataIO.ts +++ b/packages/typegpu/src/data/dataIO.ts @@ -11,6 +11,7 @@ import type { AnyConcreteData, AnyData, Disarray, LooseDecorated, Unstruct } fro import { mat2x2f, mat3x3f, mat4x4f } from './matrix.ts'; import { offsetsForProps } from './offsets.ts'; import { sizeOf } from './sizeOf.ts'; +import { strideOf } from './schemaMemoryLayout.ts'; import { vec2f, vec2h, @@ -29,7 +30,6 @@ import * as wgsl from './wgslTypes.ts'; import type { BufferWriteOptions } from '../core/buffer/buffer.ts'; import { getCompiledWriter } from './compiledIO.ts'; import { getName } from '../shared/meta.ts'; -import { roundUp } from '../mathUtils.ts'; import { logger } from '../tgpuLogger.ts'; import { readFloat16, writeFloat16 } from './float16Conversion.ts'; @@ -790,8 +790,7 @@ export function calculateOffsets( let naturalSize: number | undefined = undefined; if (wgsl.isWgslArray(schema) && Array.isArray(data)) { const arrayData = data as unknown[]; - naturalSize = - arrayData.length * roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType)); + naturalSize = arrayData.length * strideOf(schema); } else if (ArrayBuffer.isView(data) || data instanceof ArrayBuffer) { naturalSize = data.byteLength; } diff --git a/packages/typegpu/src/data/offsetUtils.ts b/packages/typegpu/src/data/offsetUtils.ts index 7ab5dd494e..a1c82f7af7 100644 --- a/packages/typegpu/src/data/offsetUtils.ts +++ b/packages/typegpu/src/data/offsetUtils.ts @@ -1,7 +1,6 @@ -import { roundUp } from '../mathUtils.ts'; -import { alignmentOf } from './alignmentOf.ts'; import { type OffsetInfo as PropOffsetInfo, offsetsForProps } from './offsets.ts'; import { sizeOf } from './sizeOf.ts'; +import { strideOf } from './schemaMemoryLayout.ts'; import { isContiguous } from './isContiguous.ts'; import { getLongestContiguousPrefix } from './getLongestContiguousPrefix.ts'; import type { @@ -233,7 +232,7 @@ function makeMatProxy(mat: MatData, parent: OffsetProxy): unknown { function makeArrayProxy(array: WgslArray, parent: OffsetProxy): unknown { const elementType = array.elementType as AnyWgslData; const elementSize = sizeOf(elementType); - const stride = roundUp(elementSize, alignmentOf(elementType)); + const stride = strideOf(array); const hasPadding = stride > elementSize; const elementLCP = getLongestContiguousPrefix(elementType); diff --git a/packages/typegpu/src/data/partialIO.ts b/packages/typegpu/src/data/partialIO.ts index 8b4278bd62..9dc9f09f18 100644 --- a/packages/typegpu/src/data/partialIO.ts +++ b/packages/typegpu/src/data/partialIO.ts @@ -1,11 +1,10 @@ import { BufferWriter, getSystemEndianness } from 'typed-binary'; -import { roundUp } from '../mathUtils.ts'; -import { alignmentOf } from './alignmentOf.ts'; import { getCompiledWriter } from './compiledIO.ts'; import { writeData } from './dataIO.ts'; import { isDisarray, isUnstruct } from './dataTypes.ts'; import { offsetsForProps } from './offsets.ts'; import { sizeOf } from './sizeOf.ts'; +import { strideOf } from './schemaMemoryLayout.ts'; import type * as wgsl from './wgslTypes.ts'; import { isWgslArray, isWgslStruct, type BaseData } from './wgslTypes.ts'; import type { InferPatch } from '../shared/repr.ts'; @@ -101,7 +100,7 @@ export function getPatchInstructions( if (isWgslArray(node) || isDisarray(node)) { const arrSchema = node as wgsl.WgslArray; - const elementSize = roundUp(sizeOf(arrSchema.elementType), alignmentOf(arrSchema)); + const elementSize = strideOf(arrSchema); const elementPadding = elementSize - sizeOf(arrSchema.elementType); if (ArrayBuffer.isView(value)) { diff --git a/packages/typegpu/src/data/schemaMemoryLayout.ts b/packages/typegpu/src/data/schemaMemoryLayout.ts index 691e4fd0c6..430d3c65c7 100644 --- a/packages/typegpu/src/data/schemaMemoryLayout.ts +++ b/packages/typegpu/src/data/schemaMemoryLayout.ts @@ -178,7 +178,7 @@ function computeMLOfWgslArray(data: WgslArray): SchemaMemoryLayout { const elementType = data.elementType; const elementMemoryLayout = computeMemoryLayout(elementType); const elementSize = elementMemoryLayout.size; - const stride = roundUp(elementSize, alignmentOf(elementType)); + const stride = strideOf(data); const hasPadding = stride > elementSize; const isContiguous = !hasPadding && elementMemoryLayout.isContiguous; @@ -199,7 +199,7 @@ function computeMLOfDisarray(data: Disarray): SchemaMemoryLayout { const elementType = data.elementType; const elementMemoryLayout = computeMemoryLayout(elementType); const elementSize = elementMemoryLayout.size; - const stride = roundUp(elementSize, customAlignmentOf(elementType)); + const stride = strideOf(data); const hasPadding = stride > elementSize; const isContiguous = !hasPadding && elementMemoryLayout.isContiguous; @@ -217,6 +217,17 @@ function computeMLOfDisarray(data: Disarray): SchemaMemoryLayout { } function computeMemoryLayout(data: BaseData): SchemaMemoryLayout { + let layout = cachedLayouts.get(data); + + if (layout === undefined) { + layout = computeMemoryLayoutUncached(data); + cachedLayouts.set(data, layout); + } + + return layout; +} + +function computeMemoryLayoutUncached(data: BaseData): SchemaMemoryLayout { const knownSize = knownSizesMap[data.type]; if (knownSize !== undefined) { @@ -268,16 +279,16 @@ function computeMemoryLayout(data: BaseData): SchemaMemoryLayout { */ const cachedLayouts = new WeakMap(); +/** + * Returns the distance (in bytes) between the starts of consecutive elements of an array. + */ +export function strideOf(schema: WgslArray | Disarray): number { + return roundUp(computeMemoryLayout(schema.elementType).size, alignmentOf(schema)); +} + export function getLayoutInfo( schema: BaseData, key: T, ): SchemaMemoryLayout[T] { - let layout = cachedLayouts.get(schema); - - if (layout === undefined) { - layout = computeMemoryLayout(schema); - cachedLayouts.set(schema, layout); - } - - return layout[key]; + return computeMemoryLayout(schema)[key]; } From 2eb6d6f97edb7895ab8e83de3d01ef2fb2d6e1b4 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Mon, 19 Oct 2026 12:43:08 +0200 Subject: [PATCH 2/4] schemaMLO cache test --- .../internal/schemaMemoryLayoutCache.test.ts | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 packages/typegpu/tests/internal/schemaMemoryLayoutCache.test.ts diff --git a/packages/typegpu/tests/internal/schemaMemoryLayoutCache.test.ts b/packages/typegpu/tests/internal/schemaMemoryLayoutCache.test.ts new file mode 100644 index 0000000000..51db128aae --- /dev/null +++ b/packages/typegpu/tests/internal/schemaMemoryLayoutCache.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, vi, it } from 'vitest'; +import { d } from 'typegpu'; +import { getCustomSize } from '../../src/data/dataTypes.ts'; +import { getLayoutInfo } from '../../src/data/schemaMemoryLayout.ts'; + +vi.mock('../../src/data/dataTypes.ts', async (importOriginal) => { + const original = await importOriginal(); + return { ...original, getCustomSize: vi.fn(original.getCustomSize) }; +}); + +describe('getLayoutInfo', () => { + it('computes the layout of each nested schema only once', () => { + const Inner = d.struct({ value: d.size(8, d.f32) }); + const Outer = d.struct({ a: Inner, b: Inner }); + const Schema = d.arrayOf(Outer, 2); + + expect(getLayoutInfo(Schema, 'size')).toBe(32); + expect(getLayoutInfo(Outer, 'size')).toBe(16); + expect(getLayoutInfo(Inner, 'size')).toBe(8); + expect(getCustomSize).toHaveBeenCalledTimes(1); + }); +}); From 5974b1aba17e62cc34357d814aa1c659df66d455 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Mon, 19 Oct 2026 12:44:32 +0200 Subject: [PATCH 3/4] small rename --- .../typegpu/src/data/schemaMemoryLayout.ts | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/packages/typegpu/src/data/schemaMemoryLayout.ts b/packages/typegpu/src/data/schemaMemoryLayout.ts index 430d3c65c7..89cb029786 100644 --- a/packages/typegpu/src/data/schemaMemoryLayout.ts +++ b/packages/typegpu/src/data/schemaMemoryLayout.ts @@ -102,7 +102,7 @@ function computeMLOfStruct(struct: WgslStruct): SchemaMemoryLayout { const hasPadding = prevSize !== size; - const propLayout = computeMemoryLayout(property); + const propLayout = getMemoryLayout(property); size += propLayout.size; if (Number.isNaN(size) && property.type !== 'array') { @@ -149,7 +149,7 @@ function computeMLOfUnstruct(data: Unstruct): SchemaMemoryLayout { isContiguous = false; } - const propLayout = computeMemoryLayout(property); + const propLayout = getMemoryLayout(property); size += propLayout.size; if (prefixEnd) { @@ -176,7 +176,7 @@ function computeMLOfUnstruct(data: Unstruct): SchemaMemoryLayout { function computeMLOfWgslArray(data: WgslArray): SchemaMemoryLayout { const elementType = data.elementType; - const elementMemoryLayout = computeMemoryLayout(elementType); + const elementMemoryLayout = getMemoryLayout(elementType); const elementSize = elementMemoryLayout.size; const stride = strideOf(data); @@ -197,7 +197,7 @@ function computeMLOfWgslArray(data: WgslArray): SchemaMemoryLayout { function computeMLOfDisarray(data: Disarray): SchemaMemoryLayout { const elementType = data.elementType; - const elementMemoryLayout = computeMemoryLayout(elementType); + const elementMemoryLayout = getMemoryLayout(elementType); const elementSize = elementMemoryLayout.size; const stride = strideOf(data); @@ -216,18 +216,18 @@ function computeMLOfDisarray(data: Disarray): SchemaMemoryLayout { return { size, isContiguous, longestContiguousPrefix }; } -function computeMemoryLayout(data: BaseData): SchemaMemoryLayout { +function getMemoryLayout(data: BaseData): SchemaMemoryLayout { let layout = cachedLayouts.get(data); if (layout === undefined) { - layout = computeMemoryLayoutUncached(data); + layout = computeMemoryLayout(data); cachedLayouts.set(data, layout); } return layout; } -function computeMemoryLayoutUncached(data: BaseData): SchemaMemoryLayout { +function computeMemoryLayout(data: BaseData): SchemaMemoryLayout { const knownSize = knownSizesMap[data.type]; if (knownSize !== undefined) { @@ -256,7 +256,7 @@ function computeMemoryLayoutUncached(data: BaseData): SchemaMemoryLayout { if (isDecorated(data) || isLooseDecorated(data)) { const size = getCustomSize(data); - const undecoratedLayout = computeMemoryLayout(undecorate(data)); + const undecoratedLayout = getMemoryLayout(undecorate(data)); if (size) { const isContiguous = size === undecoratedLayout.size && undecoratedLayout.isContiguous; @@ -267,7 +267,7 @@ function computeMemoryLayoutUncached(data: BaseData): SchemaMemoryLayout { longestContiguousPrefix: undecoratedLayout.longestContiguousPrefix, }; } - return computeMemoryLayout(data.inner); + return getMemoryLayout(data.inner); } throw new Error(`Cannot determine memory layout of data: ${data}`); @@ -283,12 +283,12 @@ const cachedLayouts = new WeakMap(); * Returns the distance (in bytes) between the starts of consecutive elements of an array. */ export function strideOf(schema: WgslArray | Disarray): number { - return roundUp(computeMemoryLayout(schema.elementType).size, alignmentOf(schema)); + return roundUp(getMemoryLayout(schema.elementType).size, alignmentOf(schema)); } export function getLayoutInfo( schema: BaseData, key: T, ): SchemaMemoryLayout[T] { - return computeMemoryLayout(schema)[key]; + return getMemoryLayout(schema)[key]; } From ec91fcd7e7517a2dae6dfe24b413b59d979105f7 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Thu, 22 Oct 2026 13:36:45 +0200 Subject: [PATCH 4/4] more stride of --- packages/typegpu/src/data/dataIO.ts | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/typegpu/src/data/dataIO.ts b/packages/typegpu/src/data/dataIO.ts index 4b70bbe724..02d35fc824 100644 --- a/packages/typegpu/src/data/dataIO.ts +++ b/packages/typegpu/src/data/dataIO.ts @@ -6,7 +6,6 @@ import { type ISerialOutput, } from 'typed-binary'; import type { Infer, InferInput, InferRecord } from '../shared/repr.ts'; -import { alignmentOf } from './alignmentOf.ts'; import type { AnyConcreteData, AnyData, Disarray, LooseDecorated, Unstruct } from './dataTypes.ts'; import { mat2x2f, mat3x3f, mat4x4f } from './matrix.ts'; import { offsetsForProps } from './offsets.ts'; @@ -80,7 +79,7 @@ function writeArrayLike( // The parent has placed us at our offset const beginning = output.currentByteOffset; - const stride = roundUp(sizeOf(schema.elementType), alignmentOf(schema)); + const stride = strideOf(schema); for (let i = 0; i < Math.min(schema.elementCount, value.length); i++) { output.seekTo(beginning + i * stride); @@ -111,7 +110,7 @@ function readArrayLike(input: ISerialInput, schema: wgsl.WgslArray | Disarray) { // The parent has placed us at our offset const beginning = input.currentByteOffset; - const stride = roundUp(sizeOf(schema.elementType), alignmentOf(schema)); + const stride = strideOf(schema); const elements: unknown[] = []; for (let i = 0; i < schema.elementCount; i++) {