wgsl: Add coverage for uniform_buffer_standard_layout (#4484)

Change layout validation and execution tests to use the same
constraints for `uniform` as `storage` when the language feature is
supported.

Author: jrprice

src/webgpu/capability_info.ts

@@ -936,6 +936,7 @@ export const kKnownWGSLLanguageFeatures = [
   'packed_4x8_integer_dot_product',
   'unrestricted_pointer_parameters',
   'pointer_composite_access',
+  'uniform_buffer_standard_layout',
 ] as const;
 
 export type WGSLLanguageFeature = (typeof kKnownWGSLLanguageFeatures)[number];

src/webgpu/shader/execution/expression/access/array/index.spec.ts

@@ -24,18 +24,18 @@ g.test('concrete_scalar')
   .desc(`Test indexing of an array of concrete scalars`)
   .params(u =>
     u
-      .combine(
-        'inputSource',
-        // 'uniform' address space requires array stride to be multiple of 16 bytes
-        allInputSources.filter(s => s !== 'uniform')
-      )
+      .combine('inputSource', allInputSources)
       .combine('elementType', ['i32', 'u32', 'f32', 'f16'] as const)
       .combine('indexType', ['i32', 'u32'] as const)
   )
   .fn(async t => {
     if (t.params.elementType === 'f16') {
       t.skipIfDeviceDoesNotHaveFeature('shader-f16');
     }
+    if (t.params.inputSource === 'uniform') {
+      // 'uniform' address space requires array stride to be multiple of 16 bytes without this language feature.
+      t.skipIfLanguageFeatureNotSupported('uniform_buffer_standard_layout');
+    }
     const elementType = Type[t.params.elementType];
     const indexType = Type[t.params.indexType];
     const cases: Case[] = [
@@ -87,15 +87,14 @@ g.test('bool')
   .specURL('https://www.w3.org/TR/WGSL/#array-access-expr')
   .desc(`Test indexing of an array of booleans`)
   .params(u =>
-    u
-      .combine(
-        'inputSource',
-        // 'uniform' address space requires array stride to be multiple of 16 bytes
-        allInputSources.filter(s => s !== 'uniform')
-      )
-      .combine('indexType', ['i32', 'u32'] as const)
+    u.combine('inputSource', allInputSources).combine('indexType', ['i32', 'u32'] as const)
   )
   .fn(async t => {
+    if (t.params.inputSource === 'uniform') {
+      // 'uniform' address space requires array stride to be multiple of 16 bytes without this language feature.
+      t.skipIfLanguageFeatureNotSupported('uniform_buffer_standard_layout');
+    }
+
     const indexType = Type[t.params.indexType];
     const cases: Case[] = [
       {
@@ -290,16 +289,16 @@ g.test('vector')
   .params(u =>
     u
       .combine('inputSource', allInputSources)
-      .expand('elementType', t =>
-        t.inputSource === 'uniform'
-          ? (['vec4i', 'vec4u', 'vec4f'] as const)
-          : (['vec4i', 'vec4u', 'vec4f', 'vec4h'] as const)
-      )
+      .expand('elementType', _ => ['vec4i', 'vec4u', 'vec4f', 'vec4h'] as const)
       .combine('indexType', ['i32', 'u32'] as const)
   )
   .fn(async t => {
     if (t.params.elementType === 'vec4h') {
       t.skipIfDeviceDoesNotHaveFeature('shader-f16');
+      if (t.params.inputSource === 'uniform') {
+        // 'uniform' address space requires array stride to be multiple of 16 bytes without this language feature.
+        t.skipIfLanguageFeatureNotSupported('uniform_buffer_standard_layout');
+      }
     }
     const elementType = Type[t.params.elementType];
     const indexType = Type[t.params.indexType];
@@ -372,6 +371,14 @@ g.test('matrix')
     if (t.params.elementType === 'f16') {
       t.skipIfDeviceDoesNotHaveFeature('shader-f16');
     }
+    if (
+      t.params.inputSource === 'uniform' &&
+      !t.hasLanguageFeature('uniform_buffer_standard_layout')
+    ) {
+      // 'uniform' address space requires array stride to be multiple of 16 bytes without this language feature.
+      const mat = Type.mat(t.params.columns, t.params.rows, Type[t.params.elementType]);
+      t.skipIf((align(mat.size, mat.alignment) & 15) !== 0);
+    }
     const elementType = Type[t.params.elementType];
     const indexType = Type[t.params.indexType];
     const matrixType = Type.mat(t.params.columns, t.params.rows, elementType);

src/webgpu/shader/execution/expression/expression.ts

@@ -102,6 +102,7 @@ function sizeAndAlignmentOf(ty: Type, source: InputSource): { size: number; alig
 
   if (ty instanceof ArrayType) {
     const out = sizeAndAlignmentOf(ty.elementType, source);
+    // MAINTENANCE_TODO(#4485): Remove this when all implementors support uniform_buffer_standard_layout.
     if (source === 'uniform') {
       out.alignment = align(out.alignment, 16);
     }
@@ -155,6 +156,7 @@ export function structLayout(
     alignment = Math.max(alignment, sizeAndAlign.alignment);
   }
 
+  // MAINTENANCE_TODO(#4485): Remove this when all implementors support uniform_buffer_standard_layout.
   if (source === 'uniform') {
     alignment = align(alignment, 16);
   }

src/webgpu/shader/execution/memory_layout.spec.ts

@@ -934,8 +934,12 @@ g.test('read_layout')
       `Skipping atomic test for non-storage address space`
     );
 
+    // If the `uniform_buffer_standard_layout` feature is supported, the `uniform` address space has
+    // the same layout constraints as `storage`.
+    const ubo_std_layout = t.hasLanguageFeature('uniform_buffer_standard_layout');
+
     t.skipIf(
-      testcase.skip_uniform === true && t.params.aspace === 'uniform',
+      !ubo_std_layout && testcase.skip_uniform === true && t.params.aspace === 'uniform',
       `Uniform requires 16 byte alignment`
     );
   })

src/webgpu/shader/types.ts

@@ -298,6 +298,7 @@ export function* generateTypes({
     if (scalarInfo.layout) {
       // Compute the layout of the array type.
       // Adjust the array element count or element type as needed.
+      // MAINTENANCE_TODO(#4485): Remove this when all implementors support uniform_buffer_standard_layout.
       if (addressSpace === 'uniform') {
         // Use a vec4 of the scalar type, to achieve a 16 byte alignment without internal padding.
         // This works for 4-byte scalar types, and does not work for f16.

src/webgpu/shader/validation/shader_io/align.spec.ts

@@ -202,6 +202,10 @@ g.test('required_alignment')
       .beginSubcases()
   )
   .fn(t => {
+    // If the `uniform_buffer_standard_layout` feature is supported, the `uniform` address space has
+    // the same layout constraints as `storage`.
+    const has_ubo_std_layout = t.hasLanguageFeature('uniform_buffer_standard_layout');
+
     // While this would fail validation, it doesn't fail for any reasons related to alignment.
     // Atomics are not allowed in uniform address space as they have to be read_write.
     if (t.params.address_space === 'uniform' && t.params.type.name.startsWith('atomic')) {
@@ -213,20 +217,20 @@ g.test('required_alignment')
       code += 'enable f16;\n';
     }
 
-    // Testing the struct case, generate the structf
+    // Testing the struct case, generate the struct
     if (t.params.type.name === 'S') {
       code += `struct S {
         a: mat4x2<f32>,          // Align 8
         b: array<vec${
-          t.params.address_space === 'storage' ? 2 : 4
+          t.params.address_space === 'storage' || has_ubo_std_layout ? 2 : 4
         }<i32>, 2>,  // Storage align 8, uniform 16
       }
       `;
     }
 
     // Alignment value listed in the spec
     const min_align =
-      t.params.address_space === 'storage'
+      t.params.address_space === 'storage' || has_ubo_std_layout
         ? `${t.params.type.storage}`
         : `${t.params.type.uniform}`;
     const align = t.params.align === 'alignment' ? min_align : t.params.align;
@@ -250,13 +254,14 @@ g.test('required_alignment')
       return vec4<f32>(.4, .2, .3, .1);
     }`;
 
-    // An array of `vec2` in uniform will not validate because, while the alignment on the array
-    // itself is fine, the `vec2` element inside the array will have the wrong alignment. Uniform
-    // requires that inner vec2 to have an align 16 which can only be done by specifying `vec4`
-    // instead.
-    const fails =
-      (t.params.address_space === 'uniform' && t.params.type.name.startsWith('array<vec2')) ||
-      align < min_align;
+    let fails = align < min_align;
+    if (!has_ubo_std_layout) {
+      // An array of `vec2` in uniform will not validate because, while the alignment on the array
+      // itself is fine, the `vec2` element inside the array will have the wrong alignment. Uniform
+      // requires that inner vec2 to have an align 16 which can only be done by specifying `vec4`
+      // instead.
+      fails ||= t.params.address_space === 'uniform' && t.params.type.name.startsWith('array<vec2');
+    }
 
     t.expectCompileResult(!fails, code);
   });

src/webgpu/shader/validation/shader_io/layout_constraints.spec.ts

@@ -525,11 +525,15 @@ ${decls}
     }
     code += `}\n`;
 
+    // If the `uniform_buffer_standard_layout` feature is supported, the `uniform` address space has
+    // the same layout constraints as `storage`.
+    const ubo_std_layout = t.hasLanguageFeature('uniform_buffer_standard_layout');
+
     const is_interface = t.params.aspace === 'uniform' || t.params.aspace === 'storage';
     const supports_atomic = t.params.aspace === 'storage' || t.params.aspace === 'workgroup';
     const expect =
       testcase.validity === true ||
-      (testcase.validity === 'non-uniform' && t.params.aspace !== 'uniform') ||
+      (testcase.validity === 'non-uniform' && (t.params.aspace !== 'uniform' || ubo_std_layout)) ||
       (testcase.validity === 'non-interface' && !is_interface) ||
       (testcase.validity === 'storage' && t.params.aspace === 'storage') ||
       (testcase.validity === 'atomic' && supports_atomic);