Add shared support for WGSL immediate validation

src/webgpu/capability_info.ts

@@ -987,6 +987,7 @@ export const kKnownWGSLLanguageFeatures = [
   'swizzle_assignment',
   'linear_indexing',
   'texture_formats_tier1',
+  'immediate_address_space',
 ] as const;
 
 export type WGSLLanguageFeature = (typeof kKnownWGSLLanguageFeatures)[number];

src/webgpu/shader/types.ts

@@ -1,5 +1,6 @@
 import { keysOf } from '../../common/util/data_tables.js';
 import { assert } from '../../common/util/util.js';
+import type { WGSLLanguageFeature } from '../capability_info.js';
 import { align } from '../util/math.js';
 
 const kDefaultArrayLength = 3;
@@ -102,7 +103,14 @@ export const kMatrixContainerTypeLayoutInfo =
   }
 } as const;
 
-export type AddressSpace = 'storage' | 'uniform' | 'private' | 'function' | 'workgroup' | 'handle';
+export type AddressSpace =
+  | 'storage'
+  | 'uniform'
+  | 'private'
+  | 'function'
+  | 'workgroup'
+  | 'immediate'
+  | 'handle';
 export type AccessMode = 'read' | 'write' | 'read_write';
 export type Scope = 'module' | 'function';
 
@@ -133,9 +141,12 @@ export type AddressSpaceInfo = {
   //   in the storage address space, must not be specified in the WGSL source.
   //   See §13.3 Address Spaces.
   spellAccessMode: Requirement;
+
+  // WGSL language feature required to use this address space, if any.
+  wgslLanguageFeature?: WGSLLanguageFeature;
 };
 
-export const kAddressSpaceInfo: Record<string, AddressSpaceInfo> = {
+export const kAddressSpaceInfo: Record<AddressSpace, AddressSpaceInfo> = {
   storage: {
     scope: 'module',
     binding: true,
@@ -171,6 +182,14 @@ export const kAddressSpaceInfo: Record<string, AddressSpaceInfo> = {
     accessModes: ['read_write'],
     spellAccessMode: 'never',
   },
+  immediate: {
+    scope: 'module',
+    binding: false,
+    spell: 'must',
+    accessModes: ['read'],
+    spellAccessMode: 'never',
+    wgslLanguageFeature: 'immediate_address_space',
+  },
   handle: {
     scope: 'module',
     binding: true,
@@ -237,8 +256,10 @@ export function* generateTypes({
   }
   const scalarType = isAtomic ? `atomic<${baseType}>` : baseType;
 
-  // Storage and uniform require host-sharable types.
-  if (addressSpace === 'storage' || addressSpace === 'uniform') {
+  // Storage, uniform, and immediate require host-shareable types.
+  const requiresHostShareable =
+    addressSpace === 'storage' || addressSpace === 'uniform' || addressSpace === 'immediate';
+  if (requiresHostShareable) {
     assert(isHostSharable(baseType), 'type ' + baseType.toString() + ' is not host sharable');
   }
 
@@ -289,6 +310,9 @@ export function* generateTypes({
 
   // Array types
   if (containerType === 'array') {
+    if (addressSpace === 'immediate') {
+      return;
+    }
     let arrayElemType: string = scalarType;
     let arrayElementCount: number = kDefaultArrayLength;
     let supportsAtomics = scalarInfo.supportsAtomics;
@@ -382,8 +406,11 @@ export function* supportedScalarTypes(p: { isAtomic: boolean; addressSpace: stri
     // Test atomics only on supported scalar types.
     if (p.isAtomic && !info.supportsAtomics) continue;
 
-    // Storage and uniform require host-sharable types.
-    const isHostShared = p.addressSpace === 'storage' || p.addressSpace === 'uniform';
+    // Storage, uniform, and immediate require host-shareable types.
+    const isHostShared =
+      p.addressSpace === 'storage' ||
+      p.addressSpace === 'uniform' ||
+      p.addressSpace === 'immediate';
     if (isHostShared && info.layout === undefined) continue;
 
     yield scalarType;

src/webgpu/shader/validation/decl/util.ts

@@ -1,3 +1,5 @@
+import { getGPU } from '../../../../common/util/navigator_gpu.js';
+import { supportsImmediateData } from '../../../../common/util/util.js';
 import {
   AccessMode,
   AddressSpace,
@@ -12,6 +14,40 @@ export type ShaderStage = 'vertex' | 'fragment' | 'compute';
 /** The list of all shader stages */
 export const kShaderStages = ['vertex', 'fragment', 'compute'] as const;
 
+export function requiredLanguageFeatureHeader(addressSpace: AddressSpace): string {
+  const feature = kAddressSpaceInfo[addressSpace].wgslLanguageFeature;
+  return feature === undefined ? '' : `requires ${feature};\n`;
+}
+
+type AddressSpaceSupportTest = {
+  readonly rec: Parameters<typeof getGPU>[0];
+  skip(message: string): never;
+  skipIfLanguageFeatureNotSupported(
+    langFeature: NonNullable<AddressSpaceInfo['wgslLanguageFeature']>
+  ): void;
+};
+
+export function skipIfImmediateDataNotSupported(t: AddressSpaceSupportTest): void {
+  if (!supportsImmediateData(getGPU(t.rec))) {
+    t.skip('Immediate data not supported');
+  }
+}
+
+export function skipIfAddressSpaceNotSupported(
+  t: AddressSpaceSupportTest,
+  addressSpace: AddressSpace
+): void {
+  if (addressSpace === 'immediate') {
+    skipIfImmediateDataNotSupported(t);
+    return;
+  }
+
+  const feature = kAddressSpaceInfo[addressSpace].wgslLanguageFeature;
+  if (feature !== undefined) {
+    t.skipIfLanguageFeatureNotSupported(feature);
+  }
+}
+
 /**
  * declareEntrypoint emits the WGSL to declare an entry point with the name, stage and body.
  * The generated function will have an appropriate return type and return statement, so that `body`
@@ -107,15 +143,16 @@ export function getVarDeclShader(
     p.explicitSpace ? p.addressSpace : '',
     p.explicitAccess ? p.accessMode : ''
   );
+  const header = requiredLanguageFeatureHeader(p.addressSpace);
 
   additionalBody = additionalBody ?? '';
 
   switch (info.scope) {
     case 'module':
-      return decl + '\n' + declareEntryPoint({ stage: p.stage, body: additionalBody });
+      return header + decl + '\n' + declareEntryPoint({ stage: p.stage, body: additionalBody });
 
     case 'function':
-      return declareEntryPoint({ stage: p.stage, body: decl + '\n' + additionalBody });
+      return header + declareEntryPoint({ stage: p.stage, body: decl + '\n' + additionalBody });
   }
 }
 

src/webgpu/shader/validation/decl/var.spec.ts

@@ -4,7 +4,7 @@
 
 import { makeTestGroup } from '../../../../common/framework/test_group.js';
 import { keysOf } from '../../../../common/util/data_tables.js';
-import { AddressSpace, kAccessModeInfo, kAddressSpaceInfo } from '../../types.js';
+import { kAccessModeInfo, kAddressSpaceInfo } from '../../types.js';
 import { ShaderValidationTest } from '../shader_validation_test.js';
 
 import {
@@ -582,9 +582,7 @@
   });
 
 // Address spaces that can hold an i32 variable.
-const kNonHandleAddressSpaces = keysOf(kAddressSpaceInfo).filter(
-  as => as !== 'handle'
-) as AddressSpace[];
+const kNonHandleAddressSpaces = keysOf(kAddressSpaceInfo).filter(as => as !== 'handle' && as !== 'immediate');
 
 g.test('explicit_access_mode')
   .desc('Validate uses of an explicit access mode on a var declaration')

src/webgpu/shader/validation/types/pointer.spec.ts

@@ -2,7 +2,7 @@
 
 import { makeTestGroup } from '../../../../common/framework/test_group.js';
 import { keysOf } from '../../../../common/util/data_tables.js';
-import { AddressSpace, kAccessModeInfo, kAddressSpaceInfo } from '../../types.js';
+import { kAccessModeInfo, kAddressSpaceInfo } from '../../types.js';
 import {
   pointerType,
   explicitSpaceExpander,
@@ -140,9 +140,7 @@
   });
 
 // Address spaces that can hold an i32 variable.
-const kNonHandleAddressSpaces = keysOf(kAddressSpaceInfo).filter(
-  as => as !== 'handle'
-) as AddressSpace[];
+const kNonHandleAddressSpaces = keysOf(kAddressSpaceInfo).filter(as => as !== 'handle' && as !== 'immediate');
 
 g.test('let_ptr_explicit_type_matches_var')
   .desc(