Fix naga panic and validate pipeline auto layout when immediate size overflows 256 bytes

Author: beicause

CHANGELOG.md

@@ -203,6 +203,7 @@ By @beholdnec in [#8505](https://github.com/gfx-rs/wgpu/pull/8505).
 - 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 immediate size when using pipeline auto layout is within device limit. By @beicause in [#9725](https://github.com/gfx-rs/wgpu/pull/9725).
 
 #### DX12
 
@@ -239,6 +240,7 @@ By @beholdnec in [#8505](https://github.com/gfx-rs/wgpu/pull/8505).
 - Fix `packSnorm2x16` and `packUnorm2x16` swap in the GLSL frontend. By @treylutton in [#9675](https://github.com/gfx-rs/wgpu/pull/9675).
 - Fixed WGSL loop-local `var` declarations without explicit initializers so they are zero-initialized each iteration. By @ruihe774 in [#9592](https://github.com/gfx-rs/wgpu/pull/9592).
 - Fix naga `immediate_slots_required` calculation, and it's moved from `FunctionInfo::immediate_slots_required` to `ModuleInfo::get_entry_point_immediate_slots_required`. By @beicause in [#9725](https://github.com/gfx-rs/wgpu/pull/9725).
+- Fix naga panic when immediate size overflow 256 bytes. By @beicause in [#9725](https://github.com/gfx-rs/wgpu/pull/9725).
 
 #### Vulkan
 

cts_runner/fail.lst

@@ -45,7 +45,6 @@ webgpu:api,validation,encoding,render_bundle:* // 81%, readonly flag normalizati
 webgpu:api,validation,image_copy,buffer_texture_copies:* // https://github.com/gfx-rs/wgpu/issues/7946
 webgpu:api,validation,layout_shader_compat:pipeline_layout_shader_exact_match:* // dx12, https://bugzilla.mozilla.org/show_bug.cgi?id=2017725
 webgpu:api,validation,non_filterable_texture:non_filterable_texture_with_filtering_sampler:* // 80%, depth textures with filtering samplers
-webgpu:api,validation,pipeline,immediates:* // 0%, https://github.com/gfx-rs/wgpu/issues/8556
 webgpu:api,validation,query_set,create:count:* // 0%, wgpu incorrectly rejects zero-count query sets
 webgpu:api,validation,queue,destroyed,* // 71%, writeBuffer/writeTexture return value, destroyed query set
 webgpu:api,validation,queue,writeBuffer:ranges:* // 0%, missing OperationError for invalid ranges

cts_runner/test.lst

@@ -272,6 +272,7 @@ webgpu:api,validation,image_copy,layout_related:required_bytes_in_copy:*
 webgpu:api,validation,image_copy,layout_related:rows_per_image_alignment:*
 webgpu:api,validation,image_copy,texture_related:*
 fails-if(dx12) webgpu:api,validation,layout_shader_compat:pipeline_layout_shader_exact_match:*
+webgpu:api,validation,pipeline,immediates:*
 webgpu:api,validation,query_set,destroy:*
 webgpu:api,validation,queue,buffer_mapped:*
 webgpu:api,validation,queue,destroyed,query_set:*

naga/src/valid/analyzer.rs

@@ -1307,7 +1307,7 @@ impl ModuleInfo {
         self.entry_point_immediate_sizes[index]
     }
 
-    pub fn get_entry_point_immediate_slots_required(&self, index: usize) -> ImmediateSlots {
+    pub fn get_entry_point_immediate_slots_required(&self, index: usize) -> Option<ImmediateSlots> {
         self.entry_point_immediate_slots_required[index]
     }
 }

naga/src/valid/immediates.rs

@@ -13,31 +13,39 @@ impl ImmediateSlots {
     }
 
     /// Compute the bitmask for a byte range [offset .. offset + size_bytes).
-    pub const fn from_range(offset: u32, size_bytes: u32) -> Self {
+    /// Return `None` if the bitmask overflows.
+    pub const fn from_range(offset: u32, size_bytes: u32) -> Option<Self> {
+        let Some(end) = offset.checked_add(size_bytes) else {
+            return None;
+        };
+        if end > u64::BITS * 4 {
+            return None;
+        }
         if size_bytes == 0 {
-            return Self(0);
+            return Some(Self(0));
         }
         let lo = offset / 4;
         let hi = (offset + size_bytes).div_ceil(4);
-        Self(u64::MAX << lo & u64::MAX >> (64 - hi))
+        Some(Self(u64::MAX << lo & u64::MAX >> (64 - hi)))
     }
 
     /// Compute the slots occupied by a type at a given byte offset,
     /// excluding padding between struct members.
+    /// Return `None` if the bitmask overflows.
     pub fn from_type(
         ty: &crate::TypeInner,
         offset: u32,
         types: &crate::UniqueArena<crate::Type>,
         gctx: crate::proc::GlobalCtx,
-    ) -> Self {
+    ) -> Option<Self> {
         match *ty {
             crate::TypeInner::Struct { ref members, .. } => {
                 let mut slots = Self::default();
                 for member in members {
                     let member_ty = &types[member.ty].inner;
-                    slots |= Self::from_type(member_ty, offset + member.offset, types, gctx);
+                    slots |= Self::from_type(member_ty, offset + member.offset, types, gctx)?;
                 }
-                slots
+                Some(slots)
             }
             _ => Self::from_range(offset, ty.size(gctx)),
         }
@@ -101,35 +109,35 @@ mod tests {
     #[test]
     fn range_single() {
         assert_eq!(
-            ImmediateSlots::from_range(0, 4),
+            ImmediateSlots::from_range(0, 4).unwrap(),
             ImmediateSlots::from_raw(0b1)
         );
         assert_eq!(
-            ImmediateSlots::from_range(4, 4),
+            ImmediateSlots::from_range(4, 4).unwrap(),
             ImmediateSlots::from_raw(0b10)
         );
         assert_eq!(
-            ImmediateSlots::from_range(8, 4),
+            ImmediateSlots::from_range(8, 4).unwrap(),
             ImmediateSlots::from_raw(0b100)
         );
     }
 
     #[test]
     fn range_vec4() {
         assert_eq!(
-            ImmediateSlots::from_range(0, 16),
+            ImmediateSlots::from_range(0, 16).unwrap(),
             ImmediateSlots::from_raw(0b1111)
         );
         assert_eq!(
-            ImmediateSlots::from_range(16, 16),
+            ImmediateSlots::from_range(16, 16).unwrap(),
             ImmediateSlots::from_raw(0b1111_0000)
         );
     }
 
     #[test]
     fn range_full_256() {
         assert_eq!(
-            ImmediateSlots::from_range(0, 256),
+            ImmediateSlots::from_range(0, 256).unwrap(),
             ImmediateSlots::from_raw(u64::MAX)
         );
     }
@@ -140,18 +148,19 @@ mod tests {
         let struct_ty = (module.types.iter().map(|ty| ty.1))
             .find(|ty| ty.name.as_deref() == Some("S"))
             .unwrap();
-        let slots = ImmediateSlots::from_type(&struct_ty.inner, 0, &module.types, module.to_ctx());
+        let slots =
+            ImmediateSlots::from_type(&struct_ty.inner, 0, &module.types, module.to_ctx()).unwrap();
         assert_eq!(slots, ImmediateSlots::from_raw(0b1111_0001));
     }
 
     #[test]
     fn range_unaligned() {
         assert_eq!(
-            ImmediateSlots::from_range(0, 3),
+            ImmediateSlots::from_range(0, 3).unwrap(),
             ImmediateSlots::from_raw(0b1)
         );
         assert_eq!(
-            ImmediateSlots::from_range(0, 5),
+            ImmediateSlots::from_range(0, 5).unwrap(),
             ImmediateSlots::from_raw(0b11)
         );
     }
@@ -161,16 +170,16 @@ mod tests {
         let required = ImmediateSlots::from_raw(0b1111_0001);
         let mut set = ImmediateSlots::default();
         assert!(!set.contains(required));
-        set |= ImmediateSlots::from_range(0, 4);
+        set |= ImmediateSlots::from_range(0, 4).unwrap();
         assert!(!set.contains(required));
-        set |= ImmediateSlots::from_range(16, 16);
+        set |= ImmediateSlots::from_range(16, 16).unwrap();
         assert!(set.contains(required));
     }
 
     #[test]
     fn difference() {
         let required = ImmediateSlots::from_raw(0b1111_0001);
-        let set = ImmediateSlots::from_range(0, 4);
+        let set = ImmediateSlots::from_range(0, 4).unwrap();
         assert_eq!(
             required.difference(set),
             ImmediateSlots::from_raw(0b1111_0000)

naga/src/valid/interface.rs

@@ -898,7 +898,7 @@ impl VaryingContext<'_> {
 pub(super) struct ValidatedEntryPoint {
     pub(super) func: FunctionInfo,
     pub(super) immediate_size: u32,
-    pub(super) immediate_slots_required: ImmediateSlots,
+    pub(super) immediate_slots_required: Option<ImmediateSlots>,
 }
 
 impl super::Validator {
@@ -1435,7 +1435,7 @@ impl super::Validator {
         }
 
         let mut immediate_size: u32 = 0;
-        let mut immediate_slots_required = ImmediateSlots::default();
+        let mut immediate_slots_required = Some(ImmediateSlots::default());
         {
             let mut used_immediates = module
                 .global_variables
@@ -1454,7 +1454,7 @@ impl super::Validator {
                 let ty = &module.types[module.global_variables[immediate].ty].inner;
                 let ctx = module.to_ctx();
                 immediate_size = ty.size(ctx);
-                immediate_slots_required |= ImmediateSlots::from_type(ty, 0, &module.types, ctx);
+                immediate_slots_required = ImmediateSlots::from_type(ty, 0, &module.types, ctx);
             }
         }
 

naga/src/valid/mod.rs

@@ -346,7 +346,7 @@ pub struct ModuleInfo {
     functions: Vec<FunctionInfo>,
     entry_points: Vec<FunctionInfo>,
     entry_point_immediate_sizes: Vec<u32>,
-    entry_point_immediate_slots_required: Vec<ImmediateSlots>,
+    entry_point_immediate_slots_required: Vec<Option<ImmediateSlots>>,
     const_expression_types: Box<[TypeResolution]>,
 }
 

wgpu-core/src/command/bundle.rs

@@ -816,7 +816,8 @@ fn set_immediates(state: &mut State, offset: u32, data: &[u8]) -> Result<(), Imm
         state.immediates.resize(end_offset, 0);
     }
     state.immediates[(offset as usize)..][..data.len()].copy_from_slice(data);
-    state.immediate_slots_set |= naga::valid::ImmediateSlots::from_range(offset, data.len() as u32);
+    state.immediate_slots_set |=
+        naga::valid::ImmediateSlots::from_range(offset, data.len() as u32).unwrap();
     state.immediate_dirty = true;
     Ok(())
 }

wgpu-core/src/command/pass.rs

@@ -246,7 +246,8 @@ where
         state.immediates.resize(end_offset, 0);
     }
     state.immediates[(offset as usize)..][..data.len()].copy_from_slice(data);
-    state.immediate_slots_set |= naga::valid::ImmediateSlots::from_range(offset, data.len() as u32);
+    state.immediate_slots_set |=
+        naga::valid::ImmediateSlots::from_range(offset, data.len() as u32).unwrap();
     state.immediates_dirty = true;
 
     Ok(())

wgpu-core/src/device/resource.rs

@@ -3862,6 +3862,12 @@ impl Device {
         mut derived_group_layouts: Box<ArrayVec<bgl::EntryMap, { hal::MAX_BIND_GROUPS }>>,
         immediate_size: u32,
     ) -> Result<Arc<binding_model::PipelineLayout>, pipeline::ImplicitLayoutError> {
+        if self.limits.max_immediate_size < immediate_size {
+            return Err(pipeline::ImplicitLayoutError::ImmediateRangeTooLarge {
+                size: immediate_size,
+                max: self.limits.max_immediate_size,
+            });
+        }
         while derived_group_layouts
             .last()
             .is_some_and(|map| map.is_empty())
@@ -4034,7 +4040,7 @@ impl Device {
             device: self.clone(),
             _shader_module: shader_module,
             late_sized_buffer_groups,
-            immediate_slots_required,
+            immediate_slots_required: immediate_slots_required.unwrap(),
             label: desc.label.to_string(),
             tracking_data: TrackingData::new(self.tracker_indices.compute_pipelines.clone()),
         };
@@ -4481,7 +4487,7 @@ impl Device {
         let mut _vertex_entry_point_name = String::new();
         let mut _task_entry_point_name = String::new();
         let mut _mesh_entry_point_name = String::new();
-        let mut immediate_slots_required = naga::valid::ImmediateSlots::default();
+        let mut immediate_slots_required = Some(naga::valid::ImmediateSlots::default());
         let mut immediate_size_required = 0u32;
         match desc.vertex {
             pipeline::RenderPipelineVertexProcessor::Vertex(ref vertex) => {
@@ -4514,7 +4520,15 @@ impl Device {
                                 stage.to_naga(),
                                 &_vertex_entry_point_name,
                             );
-                        immediate_slots_required |= _immediate_slots_required;
+                        if let (
+                            Some(mut immediate_slots_required),
+                            Some(_immediate_slots_required),
+                        ) = (immediate_slots_required, _immediate_slots_required)
+                        {
+                            immediate_slots_required |= _immediate_slots_required;
+                        } else {
+                            immediate_slots_required = None;
+                        }
                         immediate_size_required = immediate_size_required.max(_immediate_size);
                         io = interface
                             .check_stage(
@@ -4565,7 +4579,15 @@ impl Device {
                                 stage.to_naga(),
                                 &_task_entry_point_name,
                             );
-                        immediate_slots_required |= _immediate_slots_required;
+                        if let (
+                            Some(mut immediate_slots_required),
+                            Some(_immediate_slots_required),
+                        ) = (immediate_slots_required, _immediate_slots_required)
+                        {
+                            immediate_slots_required |= _immediate_slots_required;
+                        } else {
+                            immediate_slots_required = None;
+                        }
                         immediate_size_required = immediate_size_required.max(_immediate_size);
                         io = interface
                             .check_stage(
@@ -4614,7 +4636,15 @@ impl Device {
                                 stage.to_naga(),
                                 &_mesh_entry_point_name,
                             );
-                        immediate_slots_required |= _immediate_slots_required;
+                        if let (
+                            Some(mut immediate_slots_required),
+                            Some(_immediate_slots_required),
+                        ) = (immediate_slots_required, _immediate_slots_required)
+                        {
+                            immediate_slots_required |= _immediate_slots_required;
+                        } else {
+                            immediate_slots_required = None;
+                        }
                         immediate_size_required = immediate_size_required.max(_immediate_size);
                         io = interface
                             .check_stage(
@@ -4673,7 +4703,13 @@ impl Device {
                             stage.to_naga(),
                             &fragment_entry_point_name,
                         );
-                    immediate_slots_required |= _immediate_slots_required;
+                    if let (Some(mut immediate_slots_required), Some(_immediate_slots_required)) =
+                        (immediate_slots_required, _immediate_slots_required)
+                    {
+                        immediate_slots_required |= _immediate_slots_required;
+                    } else {
+                        immediate_slots_required = None;
+                    }
                     immediate_size_required = immediate_size_required.max(_immediate_size);
                     io = interface
                         .check_stage(
@@ -4905,7 +4941,7 @@ impl Device {
             strip_index_format: desc.primitive.strip_index_format,
             vertex_steps,
             late_sized_buffer_groups,
-            immediate_slots_required,
+            immediate_slots_required: immediate_slots_required.unwrap(),
             label: desc.label.to_string(),
             tracking_data: TrackingData::new(self.tracker_indices.render_pipelines.clone()),
             is_mesh,