Add shader validation tests on subgroup-size-control - Part II

src/webgpu/shader/validation/extension/subgroup_size_control.spec.ts

@@ -3,6 +3,8 @@ Validation tests for the subgroup_size_control extension
 `;
 
 import { makeTestGroup } from '../../../../common/framework/test_group.js';
+import { keysOf } from '../../../../common/util/data_tables.js';
+import { isPowerOfTwo } from '../../../util/math.js';
 import { UniqueFeaturesAndLimitsShaderValidationTest } from '../shader_validation_test.js';
 
 export const g = makeTestGroup(UniqueFeaturesAndLimitsShaderValidationTest);
@@ -101,3 +103,331 @@ g.test('subgroup_size_attribute_only_valid_in_compute_stage')
 
     t.expectCompileResult(stage === 'compute', kStageShaders[stage]);
   });
+
+const kSubgroupSizeValueCases = {
+  literal_abstract_int: { expr: '4', decl: '', pass: true },
+  literal_u32: { expr: '4u', decl: '', pass: true },
+  literal_i32: { expr: '4i', decl: '', pass: true },
+  const_i32: { expr: 'k_i32', decl: 'const k_i32: i32 = 4;', pass: true },
+  const_u32: { expr: 'k_u32', decl: 'const k_u32: u32 = 4;', pass: true },
+  const_expr_abstract: { expr: '2 + 2', decl: '', pass: true },
+  const_expr_named: { expr: 'k + 1', decl: 'const k = 3;', pass: true },
+  override_i32: { expr: 'o_i32', decl: 'override o_i32: i32 = 4;', pass: true },
+  override_u32: { expr: 'o_u32', decl: 'override o_u32: u32 = 4;', pass: true },
+  override_expr: { expr: 'o + 1', decl: 'override o: u32 = 3;', pass: true },
+  literal_f32: { expr: '4.0f', decl: '', pass: false },
+  literal_abstract_float: { expr: '4.0', decl: '', pass: false },
+  literal_bool: { expr: 'true', decl: '', pass: false },
+  const_f32: { expr: 'k_f32', decl: 'const k_f32: f32 = 4.0;', pass: false },
+  const_bool: { expr: 'k_bool', decl: 'const k_bool: bool = true;', pass: false },
+  override_f32: { expr: 'o_f32', decl: 'override o_f32: f32 = 4.0;', pass: false },
+  let_u32: { expr: 'r', decl: 'fn dummy() -> u32 { let r: u32 = 4; return r; }', pass: false },
+  var_u32: { expr: 'v', decl: 'var<private> v: u32 = 4;', pass: false },
+};
+
+g.test('subgroup_size_value_must_be_const_or_override_i32_u32')
+  .desc(
+    `Checks that the value of @subgroup_size must be a constant expression or an override
+     expression that resolves to an i32 or a u32.`
+  )
+  .params(u => u.combine('case', keysOf(kSubgroupSizeValueCases)))
+  .beforeAllSubcases(t => {
+    t.selectDeviceOrSkipTestCase({
+      requiredFeatures: ['subgroup-size-control' as GPUFeatureName],
+    });
+  })
+  .fn(t => {
+    const { expr, decl, pass } = kSubgroupSizeValueCases[t.params.case];
+
+    t.expectCompileResult(
+      pass,
+      `
+        enable subgroups;
+        enable subgroup_size_control;
+        ${decl}
+        @compute @workgroup_size(4) @subgroup_size(${expr})
+        fn main() {}
+      `
+    );
+  });
+
+g.test('subgroup_size_constant_value_must_be_power_of_2')
+  .desc(
+    `Checks that when @subgroup_size is a constant expression, it is a shader creation error if
+    the value is not a power of 2.`
+  )
+  .params(u =>
+    u.combine('size', [0, 1, 2, 3, 4, 5, 8, 9, 15, 16, 32, 33, 48, 64, 65, 100, 128, 256] as const)
+  )
+  .beforeAllSubcases(t => {
+    t.selectDeviceOrSkipTestCase({
+      requiredFeatures: ['subgroup-size-control' as GPUFeatureName],
+    });
+  })
+  .fn(t => {
+    const { size } = t.params;
+
+    t.expectCompileResult(
+      isPowerOfTwo(size),
+      `
+        enable subgroups;
+        enable subgroup_size_control;
+        @compute @workgroup_size(${size}) @subgroup_size(${size})
+        fn main() {}
+      `
+    );
+  });
+
+g.test('subgroup_size_override_must_be_power_of_2_at_pipeline_creation')
+  .desc(
+    `Checks that when @subgroup_size is an override expression, it is a pipeline creation error
+     if the override value resolves to a value that is not a power of 2.`
+  )
+  .params(u => u.combine('size', [3, 5, 7, 15, 31, 63, 127] as const))
+  .beforeAllSubcases(t => {
+    t.selectDeviceOrSkipTestCase({
+      requiredFeatures: ['subgroup-size-control' as GPUFeatureName],
+    });
+  })
+  .fn(t => {
+    const { size } = t.params;
+
+    const code = `
+      enable subgroups;
+      enable subgroup_size_control;
+      override S: u32;
+      @compute @workgroup_size(S) @subgroup_size(S)
+      fn main() {}
+    `;
+
+    const shaderModule = t.device.createShaderModule({ code });
+
+    t.expectGPUError(
+      'validation',
+      () => {
+        t.device.createComputePipeline({
+          layout: 'auto',
+          compute: {
+            module: shaderModule,
+            entryPoint: 'main',
+            constants: { S: size },
+          },
+        });
+      },
+      true
+    );
+  });
+
+/**
+ * Returns all valid subgroup sizes for the given adapter info, i.e. all power-of-two values
+ * between subgroupMinSize and subgroupMaxSize inclusive.
+ */
+async function getValidSubgroupSizes(device: GPUDevice): Promise<number[]> {
+  interface SubgroupProperties extends GPUAdapterInfo {
+    subgroupMinSize: number;
+    subgroupMaxSize: number;
+  }
+  const { subgroupMinSize, subgroupMaxSize } = device.adapterInfo as SubgroupProperties;
+
+  const sizes: number[] = [];
+  for (let subgroupSize = subgroupMinSize; subgroupSize <= subgroupMaxSize; subgroupSize *= 2) {
+    const wgsl = `
+enable subgroups;
+enable subgroup_size_control;
+
+@compute @workgroup_size(${subgroupSize}, 1, 1) @subgroup_size(${subgroupSize})
+fn main(@builtin(local_invocation_index) lid : u32) {
+}`;
+    device.pushErrorScope('validation');
+    const module = device.createShaderModule({ code: wgsl });
+    device.createComputePipeline({
+      layout: 'auto',
+      compute: { module, entryPoint: 'main' },
+    });
+    const error = await device.popErrorScope();
+    if (error) {
+      continue;
+    }
+    sizes.push(subgroupSize);
+  }
+  return sizes;
+}
+
+g.test('subgroup_size_override_valid_values_no_error')
+  .desc(
+    `Checks that when @subgroup_size is an override expression and the override value resolves
+     to a valid subgroup size (a power of 2 between subgroupMinSize and subgroupMaxSize), pipeline
+     creation succeeds without error.`
+  )
+  .beforeAllSubcases(t => {
+    t.selectDeviceOrSkipTestCase({
+      requiredFeatures: ['subgroup-size-control' as GPUFeatureName],
+    });
+  })
+  .fn(async t => {
+    const validSubgroupSizes = await getValidSubgroupSizes(t.device);
+    t.expect(validSubgroupSizes.length > 0, 'Expected at least one valid subgroup size');
+
+    const code = `
+      enable subgroups;
+      enable subgroup_size_control;
+      override S: u32;
+      @compute @workgroup_size(S) @subgroup_size(S)
+      fn main() {}
+    `;
+
+    const shaderModule = t.device.createShaderModule({ code });
+
+    for (const subgroupSize of validSubgroupSizes) {
+      t.expectGPUError(
+        'validation',
+        () => {
+          t.device.createComputePipeline({
+            layout: 'auto',
+            compute: {
+              module: shaderModule,
+              entryPoint: 'main',
+              constants: { S: subgroupSize },
+            },
+          });
+        },
+        false
+      );
+    }
+  });
+
+g.test('workgroup_size_x_must_be_multiple_of_subgroup_size_at_pipeline_creation')
+  .desc(
+    `Checks that a pipeline-creation error results if the x-dimension of the entry point's
+     workgroup_size is not a multiple of the subgroup_size value. Tests all combinations of
+     constant and override expressions for both workgroup_size and subgroup_size.`
+  )
+  .params(u =>
+    u
+      .combine('workgroupSizeIsOverride', [false, true] as const)
+      .combine('subgroupSizeIsOverride', [false, true] as const)
+      .combine('offset', [0, 1, -1] as const)
+      .combine('multiplier', [1, 2, 3] as const)
+  )
+  .beforeAllSubcases(t => {
+    t.selectDeviceOrSkipTestCase({
+      requiredFeatures: ['subgroup-size-control' as GPUFeatureName],
+    });
+  })
+  .fn(async t => {
+    const { workgroupSizeIsOverride, subgroupSizeIsOverride, offset, multiplier } = t.params;
+    const validSubgroupSizes = await getValidSubgroupSizes(t.device);
+    t.expect(validSubgroupSizes.length > 0, 'Expected at least one valid subgroup size');
+
+    for (const subgroupSize of validSubgroupSizes) {
+      const workgroupSizeX = subgroupSize * multiplier + offset;
+      if (workgroupSizeX <= 0) continue;
+
+      const isMultiple = workgroupSizeX % subgroupSize === 0;
+
+      if (!workgroupSizeIsOverride && !subgroupSizeIsOverride) {
+        // Both are constants
+        const code = `
+            enable subgroups;
+            enable subgroup_size_control;
+            @compute @workgroup_size(${workgroupSizeX}) @subgroup_size(${subgroupSize})
+            fn main() {}
+        `;
+
+        const shaderModule = t.device.createShaderModule({ code });
+
+        t.expectGPUError(
+          'validation',
+          () => {
+            t.device.createComputePipeline({
+              layout: 'auto',
+              compute: {
+                module: shaderModule,
+                entryPoint: 'main',
+              },
+            });
+          },
+          !isMultiple
+        );
+      } else if (workgroupSizeIsOverride && subgroupSizeIsOverride) {
+        // Both are overrides
+        const code = `
+          enable subgroups;
+          enable subgroup_size_control;
+          override S: u32;
+          override W: u32;
+          @compute @workgroup_size(W) @subgroup_size(S)
+          fn main() {}
+        `;
+
+        const shaderModule = t.device.createShaderModule({ code });
+
+        t.expectGPUError(
+          'validation',
+          () => {
+            t.device.createComputePipeline({
+              layout: 'auto',
+              compute: {
+                module: shaderModule,
+                entryPoint: 'main',
+                constants: { S: subgroupSize, W: workgroupSizeX },
+              },
+            });
+          },
+          !isMultiple
+        );
+      } else if (workgroupSizeIsOverride && !subgroupSizeIsOverride) {
+        // workgroup_size is override, subgroup_size is constant
+        const code = `
+          enable subgroups;
+          enable subgroup_size_control;
+          override W: u32;
+          @compute @workgroup_size(W) @subgroup_size(${subgroupSize})
+          fn main() {}
+        `;
+
+        const shaderModule = t.device.createShaderModule({ code });
+
+        t.expectGPUError(
+          'validation',
+          () => {
+            t.device.createComputePipeline({
+              layout: 'auto',
+              compute: {
+                module: shaderModule,
+                entryPoint: 'main',
+                constants: { W: workgroupSizeX },
+              },
+            });
+          },
+          !isMultiple
+        );
+      } else {
+        // workgroup_size is constant, subgroup_size is override
+        const code = `
+          enable subgroups;
+          enable subgroup_size_control;
+          override S: u32;
+          @compute @workgroup_size(${workgroupSizeX}) @subgroup_size(S)
+          fn main() {}
+        `;
+
+        const shaderModule = t.device.createShaderModule({ code });
+
+        t.expectGPUError(
+          'validation',
+          () => {
+            t.device.createComputePipeline({
+              layout: 'auto',
+              compute: {
+                module: shaderModule,
+                entryPoint: 'main',
+                constants: { S: subgroupSize },
+              },
+            });
+          },
+          !isMultiple
+        );
+      }
+    }
+  });