Validate Limits::max_compute_workgroup_storage_size

CHANGELOG.md

@@ -234,6 +234,7 @@ By @stuartparmenter in [#9658](https://github.com/gfx-rs/wgpu/pull/9658).
 - Add clip distances validation for `maxInterStageShaderVariables`. By @ErichDonGubler in [#8762](https://github.com/gfx-rs/wgpu/pull/8762). This may break some existing programs, but it compiles with the WebGPU spec.
 - Bring immediates in line with webgpu spec. By @atlv24 in [#9280](https://github.com/gfx-rs/wgpu/pull/9280).
 - Validate `LoadOp` and `StoreOp` are `None` for attachments without corresponding depth or stencil aspect. By @beicause in [#9567](https://github.com/gfx-rs/wgpu/pull/9567).
+- Validate `Limits::max_compute_workgroup_storage_size`. By @ErichDonGubler in [#9732](https://github.com/gfx-rs/wgpu/pull/9732).
 
 #### DX12
 

cts_runner/fail.lst

@@ -22,9 +22,6 @@ webgpu:api,validation,buffer,create:new_usages,* // https://github.com/denoland/
 webgpu:api,validation,buffer,mapping:* // crash
 webgpu:api,validation,capability_checks,features,* // 21%, TypeError not thrown for missing features (needs deno_webgpu fix)
 webgpu:api,validation,capability_checks,limits,* // 60%, workgroup storage size validation unimplemented; interstage vars counting; bind group timing on dx12
-webgpu:api,validation,compute_pipeline:limits,workgroup_storage_size:* // 0%
-webgpu:api,validation,compute_pipeline:overrides,workgroup_size,limits,workgroup_storage_size:* // 0%, missing workgroup storage size validation
-webgpu:api,validation,compute_pipeline:overrides,workgroup_size,limits:* // 0%
 webgpu:api,validation,createBindGroup:buffer,resource_state:* // 0%, https://github.com/gfx-rs/wgpu/issues/7881
 webgpu:api,validation,createBindGroup:external_texture,* // 0%, no external external texture in deno
 webgpu:api,validation,createBindGroup:texture,resource_state:* // crash, https://github.com/gfx-rs/wgpu/issues/7881

cts_runner/test.lst

@@ -92,12 +92,15 @@ webgpu:api,validation,capability_checks,limits,minUniformBufferOffsetAlignment:v
 webgpu:api,validation,compute_pipeline:basic:*
 webgpu:api,validation,compute_pipeline:limits,invocations_per_workgroup,each_component:*
 webgpu:api,validation,compute_pipeline:limits,invocations_per_workgroup:*
+webgpu:api,validation,compute_pipeline:limits,workgroup_storage_size:*
 webgpu:api,validation,compute_pipeline:overrides,entry_point,validation_error:*
 webgpu:api,validation,compute_pipeline:overrides,identifier:*
 webgpu:api,validation,compute_pipeline:overrides,uninitialized:*
 webgpu:api,validation,compute_pipeline:overrides,value,type_error:*
 webgpu:api,validation,compute_pipeline:overrides,value,validation_error,f16:*
 webgpu:api,validation,compute_pipeline:overrides,value,validation_error:*
+webgpu:api,validation,compute_pipeline:overrides,workgroup_size,limits,workgroup_storage_size:*
+webgpu:api,validation,compute_pipeline:overrides,workgroup_size,limits:*
 webgpu:api,validation,compute_pipeline:overrides,workgroup_size:*
 webgpu:api,validation,compute_pipeline:pipeline_layout,device_mismatch:*
 webgpu:api,validation,compute_pipeline:resource_compatibility:*

naga/src/front/wgsl/lower/template_list.rs

@@ -113,6 +113,7 @@ impl<'iter, 'source> TemplateListIter<'iter, 'source> {
         let (enumerant, span) = ctx.enumerant(expr)?;
         conv::map_address_space(enumerant, span, &ctx.enable_extensions)
     }
+
     pub fn maybe_address_space(
         &mut self,
         ctx: &ExpressionContext<'source, '_, '_>,

wgpu-core/src/validation.rs

@@ -290,6 +290,7 @@ struct EntryPoint {
     spec_constants: Vec<SpecializationConstant>,
     sampling_pairs: FastHashSet<(naga::Handle<Resource>, naga::Handle<Resource>)>,
     workgroup_size: [u32; 3],
+    workgroup_storage_size: u32,
     dual_source_blending: bool,
     task_payload_size: Option<u32>,
     mesh_info: Option<EntryPointMeshInfo>,
@@ -418,6 +419,15 @@ impl WebGpuError for InputError {
 #[derive(Clone, Debug, Error)]
 #[non_exhaustive]
 pub enum StageError {
+    #[error(
+        "Workgroup storage size {used} must be less than or equal to \
+        `Limits::{limit_desc}` ({limit})"
+    )]
+    WorkgroupStorageSizeLimitExceeded {
+        used: u32,
+        limit: u32,
+        limit_desc: &'static str,
+    },
     #[error(transparent)]
     InvalidWorkgroupSize(#[from] InvalidWorkgroupSizeError),
     #[error("Unable to find entry point '{0}'")]
@@ -549,7 +559,8 @@ impl WebGpuError for StageError {
                 var: _,
                 error,
             } => error.webgpu_error_type(),
-            Self::InvalidWorkgroupSize { .. }
+            Self::WorkgroupStorageSizeLimitExceeded { .. }
+            | Self::InvalidWorkgroupSize { .. }
             | Self::MissingEntryPoint(..)
             | Self::NoEntryPointFound
             | Self::MultipleEntryPointsFound
@@ -1214,46 +1225,51 @@ impl Interface {
     pub fn new(module: &naga::Module, info: &naga::valid::ModuleInfo, limits: wgt::Limits) -> Self {
         let mut resources = naga::Arena::new();
         let mut resource_mapping = FastHashMap::default();
+        let mut workspace_storage_size_per_var = FastHashMap::<&str, _>::default();
         for (var_handle, var) in module.global_variables.iter() {
-            let bind = match var.binding {
-                Some(br) => br,
-                _ => continue,
-            };
-            let naga_ty = &module.types[var.ty].inner;
+            if let Some(bind) = var.binding {
+                let naga_ty = &module.types[var.ty].inner;
 
-            let inner_ty = match *naga_ty {
-                naga::TypeInner::BindingArray { base, .. } => &module.types[base].inner,
-                ref ty => ty,
-            };
+                let inner_ty = match *naga_ty {
+                    naga::TypeInner::BindingArray { base, .. } => &module.types[base].inner,
+                    ref ty => ty,
+                };
 
-            let ty = match *inner_ty {
-                naga::TypeInner::Image {
-                    dim,
-                    arrayed,
-                    class,
-                } => ResourceType::Texture {
-                    dim,
-                    arrayed,
-                    class,
-                },
-                naga::TypeInner::Sampler { comparison } => ResourceType::Sampler { comparison },
-                naga::TypeInner::AccelerationStructure { vertex_return } => {
-                    ResourceType::AccelerationStructure { vertex_return }
-                }
-                ref other => ResourceType::Buffer {
-                    size: wgt::BufferSize::new(other.size(module.to_ctx()) as u64).unwrap(),
-                },
-            };
-            let handle = resources.append(
-                Resource {
-                    name: var.name.clone(),
-                    bind,
-                    ty,
-                    class: var.space,
-                },
-                Default::default(),
-            );
-            resource_mapping.insert(var_handle, handle);
+                let ty = match *inner_ty {
+                    naga::TypeInner::Image {
+                        dim,
+                        arrayed,
+                        class,
+                    } => ResourceType::Texture {
+                        dim,
+                        arrayed,
+                        class,
+                    },
+                    naga::TypeInner::Sampler { comparison } => ResourceType::Sampler { comparison },
+                    naga::TypeInner::AccelerationStructure { vertex_return } => {
+                        ResourceType::AccelerationStructure { vertex_return }
+                    }
+                    ref other => ResourceType::Buffer {
+                        size: wgt::BufferSize::new(other.size(module.to_ctx()) as u64).unwrap(),
+                    },
+                };
+                let handle = resources.append(
+                    Resource {
+                        name: var.name.clone(),
+                        bind,
+                        ty,
+                        class: var.space,
+                    },
+                    Default::default(),
+                );
+                resource_mapping.insert(var_handle, handle);
+            }
+            if var.space == naga::AddressSpace::WorkGroup {
+                workspace_storage_size_per_var.insert(
+                    var.name.as_deref().unwrap(),
+                    module.types[var.ty].inner.size(module.to_ctx()),
+                );
+            }
         }
 
         let immediate_size = naga::valid::ImmediateSlots::size_for_module(module);
@@ -1275,12 +1291,23 @@ impl Interface {
                 );
             }
 
+            let mut workgroup_storage_size = 0u32;
             for (var_handle, var) in module.global_variables.iter() {
                 let usage = info[var_handle];
-                if !usage.is_empty() && var.binding.is_some() {
-                    ep.resources.push(resource_mapping[&var_handle]);
+                if !usage.is_empty() {
+                    if var.binding.is_some() {
+                        ep.resources.push(resource_mapping[&var_handle]);
+                    }
+                    if var.space == naga::AddressSpace::WorkGroup {
+                        workgroup_storage_size = workgroup_storage_size
+                            .checked_add(
+                                workspace_storage_size_per_var[var.name.as_deref().unwrap()],
+                            )
+                            .unwrap();
+                    }
                 }
             }
+            ep.workgroup_storage_size = workgroup_storage_size;
 
             for key in info.sampling_set.iter() {
                 ep.sampling_pairs
@@ -1536,6 +1563,16 @@ impl Interface {
                 },
                 _ => unreachable!(),
             };
+
+            let workgroup_storage_used = entry_point.workgroup_storage_size.next_multiple_of(16);
+            if workgroup_storage_used > self.limits.max_compute_workgroup_storage_size {
+                return Err(StageError::WorkgroupStorageSizeLimitExceeded {
+                    used: workgroup_storage_used,
+                    limit: self.limits.max_compute_workgroup_storage_size,
+                    limit_desc: "max_compute_workgroup_storage_size",
+                });
+            }
+
             let total = workgroup_size_check.check_and_compute_total_invocations()?;
             if total == 0 {
                 return Err(StageError::InvalidWorkgroupSize(