[JS/WebGPU] Creating devices with subgroup features enabled if possible

js/web/lib/wasm/jsep/backend-webgpu.ts

@@ -243,13 +243,21 @@ export class WebGpuBackend {
       requiredFeatures,
     };
 
-    if (adapter.features.has('chromium-experimental-timestamp-query-inside-passes')) {
-      requiredFeatures.push('chromium-experimental-timestamp-query-inside-passes' as GPUFeatureName);
-    } else if (adapter.features.has('timestamp-query')) {
-      requiredFeatures.push('timestamp-query');
-    }
-    if (adapter.features.has('shader-f16')) {
-      requiredFeatures.push('shader-f16');
+    // Try requiring WebGPU features
+    const tryRequireFeature = (feature: GPUFeatureName) =>
+      adapter.features.has(feature) && requiredFeatures.push(feature) && true;
+    // Try chromium-experimental-timestamp-query-inside-passes and fallback to timestamp-query
+    if (!tryRequireFeature('chromium-experimental-timestamp-query-inside-passes' as GPUFeatureName)) {
+      tryRequireFeature('timestamp-query');
+    }
+    tryRequireFeature('shader-f16');
+    // Try subgroups, and fallback to chromium-experiental-subgroups
+    if (tryRequireFeature('subgroups' as GPUFeatureName)) {
+      // If subgroups feature is available, also try subgroups-f16
+      tryRequireFeature('subgroups-f16' as GPUFeatureName);
+    } else {
+      // Try fallback to chromium-experimental-subgroups, which also enables f16 for subgroups if available.
+      tryRequireFeature('chromium-experimental-subgroups' as GPUFeatureName);
     }
 
     this.device = await adapter.requestDevice(deviceDescriptor);

js/web/lib/wasm/jsep/init.ts

@@ -113,6 +113,36 @@ class ComputeContextImpl implements ComputeContext {
     return this.backend.device.limits.maxComputeWorkgroupStorageSize;
   }
 
+  isSubgroupsSupported(): boolean {
+    return (
+      this.backend.device.features.has('subgroups' as GPUFeatureName) ||
+      this.backend.device.features.has('chromium-experimental-subgroups' as GPUFeatureName)
+    );
+  }
+
+  isSubgroupsF16Supported(): boolean {
+    return (
+      this.backend.device.features.has('subgroups-f16' as GPUFeatureName) ||
+      (this.backend.device.features.has('chromium-experimental-subgroups' as GPUFeatureName) &&
+        this.backend.device.features.has('shader-f16'))
+    );
+  }
+
+  getSubgroupSizeRange(): [number, number] | undefined {
+    // Currently subgroups feature is still experimental and size attributes are not in the WebGPU IDL, so we have to
+    // workaround the IDL type checks.
+    // TODO: clean this after subgroups feature is sattled in IDL.
+    const deviceSubgroupsLimits = this.backend.device.limits as { minSubgroupSize?: number; maxSubgroupSize?: number };
+    if (
+      !this.isSubgroupsSupported() ||
+      !deviceSubgroupsLimits.minSubgroupSize ||
+      !deviceSubgroupsLimits.maxSubgroupSize
+    ) {
+      return undefined;
+    }
+    return [deviceSubgroupsLimits.minSubgroupSize, deviceSubgroupsLimits.maxSubgroupSize];
+  }
+
   compute(program: ProgramInfo, inputsOutputsMapping?: ComputeContextInputsOutputsMapping): TensorView[] {
     // prepare inputs. inputs should always be valid data.
     const mappedInputs =

js/web/lib/wasm/jsep/webgpu/program-manager.ts

@@ -93,13 +93,24 @@ export class ProgramManager {
   build(programInfo: ProgramInfo, normalizedDispatchGroupSize: [number, number, number]): Artifact {
     TRACE_FUNC_BEGIN(programInfo.name);
     const device = this.backend.device;
-    const extensions: string[] = [];
-    if (device.features.has('shader-f16')) {
-      extensions.push('enable f16;');
-    }
+    const enableDirectives: string[] = [];
+
+    // Enable WGSL extensions based on available WebGPU features
+    const extensionsInfo: Array<{ feature: GPUFeatureName; extension: string }> = [
+      { feature: 'shader-f16', extension: 'f16' },
+      { feature: 'subgroups' as GPUFeatureName, extension: 'subgroups' },
+      { feature: 'subgroups-f16' as GPUFeatureName, extension: 'subgroups_f16' },
+      { feature: 'chromium-experimental-subgroups' as GPUFeatureName, extension: 'chromium_experimental_subgroups' },
+    ];
+    extensionsInfo.forEach((info) => {
+      if (device.features.has(info.feature)) {
+        enableDirectives.push(`enable ${info.extension};`);
+      }
+    });
+
     const shaderHelper = createShaderHelper(normalizedDispatchGroupSize, this.backend.device.limits);
     const userCode = programInfo.getShaderSource(shaderHelper);
-    const code = `${extensions.join('\n')}\n${shaderHelper.additionalImplementations}\n${userCode}`;
+    const code = `${enableDirectives.join('\n')}\n${shaderHelper.additionalImplementations}\n${userCode}`;
     const shaderModule = device.createShaderModule({ code, label: programInfo.name });
     LOG_DEBUG('verbose', () => `[WebGPU] ${programInfo.name} shader code: ${code}`);
 

js/web/lib/wasm/jsep/webgpu/types.ts

@@ -189,6 +189,9 @@ export interface ComputeContext {
   output(index: number, dims: readonly number[]): number;
   getMaxComputeWorkgroupSizes(): [number, number, number];
   getMaxComputeWorkgroupStoragesize(): number;
+  isSubgroupsSupported(): boolean;
+  isSubgroupsF16Supported(): boolean;
+  getSubgroupSizeRange(): [number, number] | undefined;
 }
 
 export type TimestampQuery = 'none' | 'inside-passes' | 'at-passes';