impr: Better d.memoryLayoutOf cache

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})`,

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',

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<TData extends WgslArray | Disarray> 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);
   }
 

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

packages/typegpu/src/data/dataIO.ts

@@ -6,11 +6,11 @@ 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';
 import { sizeOf } from './sizeOf.ts';
+import { strideOf } from './schemaMemoryLayout.ts';
 import {
   vec2f,
   vec2h,
@@ -29,7 +29,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';
 
@@ -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++) {
@@ -790,8 +789,7 @@ export function calculateOffsets<T extends wgsl.BaseData>(
   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;
   }

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);
 

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<TData extends wgsl.BaseData>(
 
     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)) {

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,9 +176,9 @@ 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 = roundUp(elementSize, alignmentOf(elementType));
+  const stride = strideOf(data);
 
   const hasPadding = stride > elementSize;
   const isContiguous = !hasPadding && elementMemoryLayout.isContiguous;
@@ -197,9 +197,9 @@ 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 = roundUp(elementSize, customAlignmentOf(elementType));
+  const stride = strideOf(data);
 
   const hasPadding = stride > elementSize;
   const isContiguous = !hasPadding && elementMemoryLayout.isContiguous;
@@ -216,6 +216,17 @@ function computeMLOfDisarray(data: Disarray): SchemaMemoryLayout {
   return { size, isContiguous, longestContiguousPrefix };
 }
 
+function getMemoryLayout(data: BaseData): SchemaMemoryLayout {
+  let layout = cachedLayouts.get(data);
+
+  if (layout === undefined) {
+    layout = computeMemoryLayout(data);
+    cachedLayouts.set(data, layout);
+  }
+
+  return layout;
+}
+
 function computeMemoryLayout(data: BaseData): SchemaMemoryLayout {
   const knownSize = knownSizesMap[data.type];
 
@@ -245,7 +256,7 @@ function computeMemoryLayout(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;
@@ -256,7 +267,7 @@ function computeMemoryLayout(data: BaseData): SchemaMemoryLayout {
         longestContiguousPrefix: undecoratedLayout.longestContiguousPrefix,
       };
     }
-    return computeMemoryLayout(data.inner);
+    return getMemoryLayout(data.inner);
   }
 
   throw new Error(`Cannot determine memory layout of data: ${data}`);
@@ -268,16 +279,16 @@ function computeMemoryLayout(data: BaseData): SchemaMemoryLayout {
  */
 const cachedLayouts = new WeakMap<BaseData, SchemaMemoryLayout>();
 
+/**
+ * Returns the distance (in bytes) between the starts of consecutive elements of an array.
+ */
+export function strideOf(schema: WgslArray | Disarray): number {
+  return roundUp(getMemoryLayout(schema.elementType).size, alignmentOf(schema));
+}
+
 export function getLayoutInfo<T extends keyof SchemaMemoryLayout>(
   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 getMemoryLayout(schema)[key];
 }

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<typeof import('../../src/data/dataTypes.ts')>();
+  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);
+  });
+});