impr: Warn when schema is not uniform aligned

apps/typegpu-docs/src/examples/react/confetti/index.tsx

@@ -127,7 +127,7 @@ function App() {
 
   const particleDataBuffer = useBuffer(d.arrayOf(ParticleData, PARTICLE_AMOUNT), {
     initial: writeRandomPositions,
-  }).$usage('storage', 'uniform', 'vertex');
+  }).$usage('storage', 'vertex');
 
   const aspectRatio = useUniform(d.f32, { initial: 1 });
   const deltaTime = useUniform(d.f32);

apps/typegpu-docs/src/examples/simulation/confetti/index.ts

@@ -48,7 +48,7 @@ const particleGeometryBuffer = root
 
 const particleDataBuffer = root
   .createBuffer(d.arrayOf(ParticleData, PARTICLE_AMOUNT))
-  .$usage('storage', 'uniform', 'vertex');
+  .$usage('storage', 'vertex');
 
 let elapsedTime = 0;
 const aspectRatio = root.createUniform(d.f32, canvas.width / canvas.height);

packages/typegpu/src/core/buffer/buffer.ts

@@ -34,6 +34,7 @@ import {
   type TgpuReadonly,
   type TgpuUniform,
 } from './bufferBinding.ts';
+import { warnIfNotUniformAligned } from '../pipeline/webgpuLimitations.ts';
 
 // ----------
 // Public API
@@ -350,6 +351,10 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
         throw new Error(`Buffer of type ${this.dataType.type} cannot be used as ${usage}`);
       }
 
+      if (usage === 'uniform') {
+        warnIfNotUniformAligned(this.dataType);
+      }
+
       this.flags |= usage === 'uniform' ? GPUBufferUsage.UNIFORM : 0;
       this.flags |= usage === 'storage' ? GPUBufferUsage.STORAGE : 0;
       this.flags |= usage === 'vertex' ? GPUBufferUsage.VERTEX : 0;

packages/typegpu/src/core/pipeline/computePipeline.ts

@@ -39,7 +39,7 @@ import type { ExperimentalTgpuRoot } from '../root/rootTypes.ts';
 import type { TgpuSlot } from '../slot/slotTypes.ts';
 
 import type { PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
-import { warnIfOverflow } from './limitsOverflow.ts';
+import { warnIfOverflow } from './webgpuLimitations.ts';
 import {
   collectBindGroupPairs,
   DISPATCH_INDIRECT_SIZE,

packages/typegpu/src/core/pipeline/renderPipeline.ts

@@ -80,7 +80,7 @@ import {
 } from './timeable.ts';
 import { nonTransferablePriorsOf } from './priors.ts';
 import { type PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
-import { warnIfOverflow } from './limitsOverflow.ts';
+import { warnIfOverflow } from './webgpuLimitations.ts';
 import {
   collectBindGroupPairs,
   collectVertexBufferPairs,

packages/typegpu/src/core/pipeline/webgpuLimitations.ts

@@ -0,0 +1,111 @@
+import { alignmentOf } from '../../data/alignmentOf.ts';
+import { memoryLayoutOf } from '../../data/offsetUtils.ts';
+import { sizeOf } from '../../data/sizeOf.ts';
+import { isWgslArray, isWgslStruct, type BaseData } from '../../data/wgslTypes.ts';
+import { invariant } from '../../errors.ts';
+import { roundUp } from '../../mathUtils.ts';
+import { getName } from '../../shared/meta.ts';
+import type { TgpuBindGroupLayout } from '../../tgpuBindGroupLayout.ts';
+import { logger } from '../../tgpuLogger.ts';
+
+/**
+ * Warns if layout exceeds supported buffer count limits.
+ */
+export function warnIfOverflow(layouts: TgpuBindGroupLayout[], limits: GPUSupportedLimits) {
+  const entries = Object.values(layouts)
+    .flatMap((layout) => Object.values(layout.entries))
+    .filter((entry) => entry !== null);
+
+  const uniform = entries.filter((entry) => 'uniform' in entry).length;
+  const storage = entries.filter((entry) => 'storage' in entry).length;
+
+  if (uniform > limits.maxUniformBuffersPerShaderStage) {
+    logger.warn(
+      'webgpu-limits-exceeded',
+      `Total number of uniform buffers (${uniform}) exceeds maxUniformBuffersPerShaderStage (${limits.maxUniformBuffersPerShaderStage}). Consider:
+1. Grouping some of the uniforms into one using 'd.struct',
+2. Increasing the limit when requesting a device or creating a root.`,
+    );
+  }
+
+  if (storage > limits.maxStorageBuffersPerShaderStage) {
+    logger.warn(
+      'webgpu-limits-exceeded',
+      `Total number of storage buffers (${storage}) exceeds maxStorageBuffersPerShaderStage (${limits.maxStorageBuffersPerShaderStage}).`,
+    );
+  }
+}
+
+function requiredAlignOf(schema: BaseData) {
+  if (isWgslStruct(schema) || isWgslArray(schema)) {
+    return roundUp(alignmentOf(schema), 16);
+  }
+  return alignmentOf(schema);
+}
+
+/**
+ * See https://www.w3.org/TR/WGSL/#address-space-layout-constraints
+ */
+export function warnIfNotUniformAligned(schema: BaseData) {
+  if (isWgslArray(schema)) {
+    warnIfNotUniformAligned(schema.elementType);
+
+    const stride = roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType));
+    if (stride % 16) {
+      logger.warn(
+        'uniform-schema-misaligned',
+        `\
+Schema '${getName(schema.elementType) ?? '<unnamed>'}' is used in an array in a uniform buffer, and its stride (${stride}) is not a multiple of 16.
+This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+To address this, put the element schema in a struct and wrap the prop in 'd.align(16, ...)', or use other schema like 'vec4f'.`,
+      );
+    }
+  }
+  if (isWgslStruct(schema)) {
+    Object.values(schema.propTypes).forEach(warnIfNotUniformAligned);
+
+    Object.entries(schema.propTypes).forEach(([key, value]) => {
+      const offset = memoryLayoutOf(schema, (schema) => schema[key]).offset;
+      const requiredAlignment = requiredAlignOf(value);
+
+      if (offset % requiredAlignment) {
+        logger.warn(
+          'uniform-schema-misaligned',
+          `\
+Schema '${getName(schema) ?? '<unnamed>'}' is used in an uniform buffer, and its property '${key}' does not meet required alignment (offset is ${offset}, required alignment is ${requiredAlignment}).
+This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+To address this, wrap the property '${key}' in 'd.align(${requiredAlignment}, ...)'.`,
+        );
+      }
+    });
+
+    const keys = Object.keys(schema.propTypes);
+    for (let i = 0; i < keys.length - 1; i++) {
+      const thisKey = keys[i];
+      const nextKey = keys[i + 1];
+      invariant(thisKey && nextKey);
+
+      const thisValue = schema.propTypes[thisKey];
+      invariant(thisValue);
+
+      if (!isWgslStruct(thisValue)) {
+        continue;
+      }
+
+      const minimumDifference = roundUp(sizeOf(thisValue), 16);
+      const thisKeyOffset = memoryLayoutOf(schema, (schema) => schema[thisKey]).offset;
+      const nextKeyOffset = memoryLayoutOf(schema, (schema) => schema[nextKey]).offset;
+      const difference = nextKeyOffset - thisKeyOffset;
+
+      if (minimumDifference > difference) {
+        logger.warn(
+          'uniform-schema-misaligned',
+          `\
+Schema '${getName(schema) ?? '<unnamed>'}' is used in an uniform buffer, and the difference between memory offsets of '${thisKey}' and '${nextKey}' props (${difference}) is less than recommended (${minimumDifference}).
+This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+To address this, wrap the '${thisKey}' prop in 'd.size(${minimumDifference}, ...)'.`,
+        );
+      }
+    }
+  }
+}

packages/typegpu/src/tgpuLogger.ts

@@ -13,6 +13,7 @@ const warningTypes = [
   'locations-mismatched',
   'log-limit-exceeded',
   'external-omitted',
+  'uniform-schema-misaligned',
 ] as const;
 type WarningType = (typeof warningTypes)[number];
 

packages/typegpu/tests/buffer.test.ts

@@ -1773,3 +1773,164 @@ describe('ValidateBufferSchema', () => {
     );
   });
 });
+
+describe('Uniform alignment', () => {
+  it('does not report legit schemas', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createUniform(d.u32);
+    root.createUniform(d.struct({ p: d.u32 }));
+    root.createUniform(d.struct({ p: d.struct({ p: d.u32 }), q: d.vec4f }));
+    root.createUniform(d.struct({ p: d.struct({ p: d.u32 }), q: d.align(16, d.u32) }));
+    root.createUniform(d.struct({ p: d.struct({ p: d.u32 }), q: d.align(32, d.u32) }));
+    root.createUniform(d.struct({ p: d.struct({ p: d.u32 }), q: d.vec3f }));
+    root.createUniform(d.struct({ p: d.size(16, d.struct({ p: d.u32 })), q: d.u32 }));
+    root.createUniform(d.struct({ p: d.size(32, d.struct({ p: d.u32 })), q: d.u32 }));
+    root.createUniform(d.arrayOf(d.vec4f, 3));
+    root.createUniform(d.arrayOf(d.vec3f, 3));
+    root.createUniform(d.arrayOf(d.struct({ p: d.vec3f }), 3));
+
+    expect(consoleWarnSpy).not.toHaveBeenCalled();
+  });
+
+  it('reports props not meeting requiredAlignOf in structs', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createUniform(d.struct({ q: d.u32, p: d.struct({ p: d.u32 }) }));
+
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema '<unnamed>' is used in an uniform buffer, and its property 'p' does not meet required alignment (offset is 4, required alignment is 16).
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, wrap the property 'p' in 'd.align(16, ...)'.",
+      ]
+    `);
+  });
+
+  it('reports unaligned props in structs', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createUniform(d.struct({ p: d.struct({ p: d.u32 }), q: d.u32 }));
+
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema '<unnamed>' is used in an uniform buffer, and the difference between memory offsets of 'p' and 'q' props (4) is less than recommended (16).
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, wrap the 'p' prop in 'd.size(16, ...)'.",
+      ]
+    `);
+  });
+
+  it('reports further unaligned props in structs', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createUniform(d.struct({ p: d.vec4f, q: d.struct({ p: d.u32 }), r: d.u32 }));
+
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema '<unnamed>' is used in an uniform buffer, and the difference between memory offsets of 'q' and 'r' props (4) is less than recommended (16).
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, wrap the 'q' prop in 'd.size(16, ...)'.",
+      ]
+    `);
+  });
+
+  it('reports unaligned props in structs of size greater than 16', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    // The struct has size 20, therefore the next member must start at 32 or later.
+    root.createUniform(
+      d.struct({ p: d.struct({ p: d.u32, q: d.u32, r: d.u32, s: d.u32, t: d.u32 }), q: d.u32 }),
+    );
+
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema '<unnamed>' is used in an uniform buffer, and the difference between memory offsets of 'p' and 'q' props (20) is less than recommended (32).
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, wrap the 'p' prop in 'd.size(32, ...)'.",
+      ]
+    `);
+  });
+
+  it('reports nested unaligned props in structs', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createUniform(
+      d.struct({ p: d.vec4f, q: d.struct({ p: d.struct({ p: d.u32 }), q: d.u32 }) }),
+    );
+
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema 'q' is used in an uniform buffer, and the difference between memory offsets of 'p' and 'q' props (4) is less than recommended (16).
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, wrap the 'p' prop in 'd.size(16, ...)'.",
+      ]
+    `);
+  });
+
+  it('reports unaligned arrays', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createUniform(d.arrayOf(d.u32, 3));
+
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema 'u32' is used in an array in a uniform buffer, and its stride (4) is not a multiple of 16.
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, put the element schema in a struct and wrap the prop in 'd.align(16, ...)', or use other schema like 'vec4f'.",
+      ]
+    `);
+  });
+
+  it('reports nested unaligned arrays', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createUniform(d.arrayOf(d.arrayOf(d.u32, 4), 4));
+
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema 'u32' is used in an array in a uniform buffer, and its stride (4) is not a multiple of 16.
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, put the element schema in a struct and wrap the prop in 'd.align(16, ...)', or use other schema like 'vec4f'.",
+      ]
+    `);
+  });
+
+  it('reports when giving usage', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createBuffer(d.arrayOf(d.u32, 2)).$usage('uniform');
+
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema 'u32' is used in an array in a uniform buffer, and its stride (4) is not a multiple of 16.
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, put the element schema in a struct and wrap the prop in 'd.align(16, ...)', or use other schema like 'vec4f'.",
+      ]
+    `);
+  });
+
+  it('does not report twice', ({ root }) => {
+    using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    root.createBuffer(d.arrayOf(d.u32, 2)).$usage('uniform').as('uniform');
+
+    expect(consoleWarnSpy).toHaveBeenCalledTimes(1);
+    expect(consoleWarnSpy.mock.calls[0]).toMatchInlineSnapshot(`
+      [
+        "⚠️ [uniform-schema-misaligned] ",
+        "Schema 'u32' is used in an array in a uniform buffer, and its stride (4) is not a multiple of 16.
+      This is not portable (see https://www.w3.org/TR/WGSL/#address-space-layout-constraints), and will break on some devices.
+      To address this, put the element schema in a struct and wrap the prop in 'd.align(16, ...)', or use other schema like 'vec4f'.",
+      ]
+    `);
+  });
+});

packages/typegpu/tests/internal/limitsOverflow.test.ts

@@ -1,7 +1,7 @@
 import { describe, expect, vi } from 'vitest';
 import { it } from 'typegpu-testing-utility';
 import { tgpu, d } from 'typegpu';
-import { warnIfOverflow } from '../../src/core/pipeline/limitsOverflow.ts';
+import { warnIfOverflow } from '../../src/core/pipeline/webgpuLimitations.ts';
 
 describe('warnIfOverflow', () => {
   const limits = {