feat: add WebGPU subgroup operations support (#8528)

Request the `subgroups` device feature, add `supportsSubgroups` property,
auto-inject `enable subgroups;` and `requires subgroup_id;` into WGSL shaders,
and expose `CAPS_SUBGROUPS` / `CAPS_SUBGROUP_ID` shader defines.

Made-with: Cursor

Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>

Author: mvaligursky

src/platform/graphics/graphics-device.js

@@ -258,13 +258,22 @@ class GraphicsDevice extends EventHandler {
      */
     supportsStorageTextureRead = false;
 
+    /**
+     * True if the device supports subgroup operations in shaders (WebGPU only). When supported,
+     * compute and fragment shaders can use WGSL subgroup builtins such as `subgroupBroadcast`,
+     * `subgroupAll`, `subgroupAny`, `subgroupAdd`, `subgroupShuffle`, etc. The `enable subgroups;`
+     * directive is automatically injected into WGSL shaders when this feature is available.
+     *
+     * @type {boolean}
+     * @readonly
+     */
+    supportsSubgroups = false;
+
     /**
      * True if the device supports the WGSL subgroup_uniformity extension, which allows
      * subgroup functionality to be considered uniform in more cases during shader compilation.
-     * When a shader uses this feature, use an `enable` directive at the top of the WGSL shader:
-     * ```wgsl
-     * enable subgroups;
-     * ```
+     * This is automatically enabled via the `enable subgroups;` directive when
+     * {@link GraphicsDevice#supportsSubgroups} is true.
      *
      * @readonly
      * @type {boolean}
@@ -273,14 +282,11 @@ class GraphicsDevice extends EventHandler {
 
     /**
      * True if the device supports the WGSL subgroup_id extension, which provides access to
-     * `subgroup_id` and `num_subgroups` built-in values in workgroups.
-     * When a shader uses this feature, use a `requires` directive at the top of the WGSL shader:
-     * ```wgsl
-     * requires subgroup_id;
-     * ```
+     * `subgroup_id` and `num_subgroups` built-in values in workgroups. The `requires subgroup_id;`
+     * directive is automatically injected into WGSL shaders when this feature is available.
      *
-     * @readonly
      * @type {boolean}
+     * @readonly
      */
     supportsSubgroupId = false;
 
@@ -655,6 +661,8 @@ class GraphicsDevice extends EventHandler {
         if (this.supportsMultiDraw) capsDefines.set('CAPS_MULTI_DRAW', '');
         if (this.supportsPrimitiveIndex) capsDefines.set('CAPS_PRIMITIVE_INDEX', '');
         if (this.supportsShaderF16) capsDefines.set('CAPS_SHADER_F16', '');
+        if (this.supportsSubgroups) capsDefines.set('CAPS_SUBGROUPS', '');
+        if (this.supportsSubgroupId) capsDefines.set('CAPS_SUBGROUP_ID', '');
 
         // Platform defines
         if (platform.desktop) capsDefines.set('PLATFORM_DESKTOP', '');

src/platform/graphics/shader-definition-utils.js

@@ -221,6 +221,12 @@ class ShaderDefinitionUtils {
         if (shaderType === 'fragment' && device.supportsPrimitiveIndex) {
             code += 'enable primitive_index;\n';
         }
+        if (device.supportsSubgroups) {
+            code += 'enable subgroups;\n';
+        }
+        if (device.supportsSubgroupId) {
+            code += 'requires subgroup_id;\n';
+        }
         return code;
     }
 

src/platform/graphics/webgpu/webgpu-graphics-device.js

@@ -327,6 +327,7 @@ class WebgpuGraphicsDevice extends GraphicsDevice {
         this.supportsTextureFormatTier2 = requireFeature('texture-format-tier2');
         this.supportsTextureFormatTier1 ||= this.supportsTextureFormatTier2;
         this.supportsPrimitiveIndex = requireFeature('primitive-index');
+        this.supportsSubgroups = requireFeature('subgroups');
         Debug.log(`WEBGPU features: ${requiredFeatures.join(', ')}`);
 
         // copy all adapter limits to the requiredLimits object - to created a device with the best feature sets available