fix(core): Use more checked arithmetic (#9357)

andyleiserson · jimblandy · commit f52e38509f33 · 2026-04-30T17:20:47.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -122,6 +122,7 @@ depth_stencil: Some(wgpu::DepthStencilState::stencil(
 
 #### General
 
+- Added new `InvalidWorkgroupSizeError`, which is now used by `DrawError::InvalidGroupSize` and `StageError::InvalidWorkgroupSize`. By @andyleiserson in [#9357](https://github.com/gfx-rs/wgpu/pull/9357).
 - Added support for cooperative load/store operations in shaders. Currently only WGSL on the input and SPIR-V, METAL, and WGSL on the output are supported. By @kvark in [#8251](https://github.com/gfx-rs/wgpu/issues/8251).
 - Added support for per-vertex attributes in fragment shaders. Currently only WGSL input is supported, and only SPIR-V or WGSL output is supported. By @atlv24 in [#8821](https://github.com/gfx-rs/wgpu/issues/8821).
 - Added support for no-perspective barycentric coordinates. By @atlv24 in [#8852](https://github.com/gfx-rs/wgpu/issues/8852).
diff --git a/wgpu-core/src/binding_model.rs b/wgpu-core/src/binding_model.rs
@@ -5,7 +5,7 @@ use alloc::{
     sync::{Arc, Weak},
     vec::Vec,
 };
-use core::{fmt, mem::ManuallyDrop, ops::Range};
+use core::{fmt, mem::ManuallyDrop, num::Saturating, ops::Range};
 
 use arrayvec::ArrayVec;
 use thiserror::Error;
@@ -363,9 +363,9 @@ impl BindingTypeMaxCountErrorKind {
 
 #[derive(Debug, Default)]
 pub(crate) struct PerStageBindingTypeCounter {
-    vertex: u32,
-    fragment: u32,
-    compute: u32,
+    vertex: Saturating<u32>,
+    fragment: Saturating<u32>,
+    compute: Saturating<u32>,
 }
 
 impl PerStageBindingTypeCounter {
@@ -393,7 +393,7 @@ impl PerStageBindingTypeCounter {
         if max_value == self.compute {
             stage |= wgt::ShaderStages::COMPUTE
         }
-        (BindingZone::Stage(stage), max_value)
+        (BindingZone::Stage(stage), max_value.0)
     }
 
     pub(crate) fn merge(&mut self, other: &Self) {
diff --git a/wgpu-core/src/command/bundle.rs b/wgpu-core/src/command/bundle.rs
@@ -116,6 +116,7 @@ use crate::{
     resource_log,
     snatch::SnatchGuard,
     track::RenderBundleScope,
+    validation::check_workgroup_sizes,
     Label, LabelHelpers,
 };
 
@@ -791,21 +792,20 @@ fn draw_mesh_tasks(
 ) -> Result<(), RenderBundleErrorInner> {
     state.is_ready(DrawCommandFamily::DrawMeshTasks)?;
 
-    let groups_size_limit = state.device.limits.max_task_mesh_workgroups_per_dimension;
-    let max_groups = state.device.limits.max_task_mesh_workgroup_total_count;
-    if group_count_x > groups_size_limit
-        || group_count_y > groups_size_limit
-        || group_count_z > groups_size_limit
-        || group_count_x * group_count_y * group_count_z > max_groups
-    {
-        return Err(RenderBundleErrorInner::Draw(DrawError::InvalidGroupSize {
-            current: [group_count_x, group_count_y, group_count_z],
-            limit: groups_size_limit,
-            max_total: max_groups,
-        }));
-    }
-
-    if group_count_x > 0 && group_count_y > 0 && group_count_z > 0 {
+    let total_count = check_workgroup_sizes(
+        &[group_count_x, group_count_y, group_count_z],
+        &[
+            state.device.limits.max_task_mesh_workgroups_per_dimension,
+            state.device.limits.max_task_mesh_workgroups_per_dimension,
+            state.device.limits.max_task_mesh_workgroups_per_dimension,
+        ],
+        "max_task_mesh_workgroups_per_dimension",
+        state.device.limits.max_task_mesh_workgroup_total_count,
+        "max_task_mesh_workgroup_total_count",
+    )
+    .map_err(|err| RenderBundleErrorInner::Draw(err.into()))?;
+
+    if total_count > 0 {
         state.flush_bindings();
         state.commands.push(ArcRenderCommand::DrawMeshTasks {
             group_count_x,
@@ -838,6 +838,7 @@ fn multi_draw_indirect(
     let vertex_limits = super::VertexLimits::new(state.vertex_buffer_sizes(), &pipeline.steps);
 
     let stride = super::get_stride_of_indirect_args(family);
+    assert!(offset <= wgt::BufferAddress::MAX - stride);
     state
         .buffer_memory_init_actions
         .extend(buffer.initialization_status.read().create_action(
diff --git a/wgpu-core/src/command/compute.rs b/wgpu-core/src/command/compute.rs
@@ -885,11 +885,11 @@ fn dispatch_indirect(
         return Err(ComputePassErrorInner::UnalignedIndirectBufferOffset(offset));
     }
 
-    let end_offset = offset + size_of::<wgt::DispatchIndirectArgs>() as u64;
-    if end_offset > buffer.size {
+    let args_size = size_of::<wgt::DispatchIndirectArgs>() as u64;
+    if buffer.size < args_size || buffer.size - args_size < offset {
         return Err(ComputePassErrorInner::IndirectBufferOverrun {
             offset,
-            end_offset,
+            end_offset: offset + args_size,
             buffer_size: buffer.size,
         });
     }
diff --git a/wgpu-core/src/command/draw.rs b/wgpu-core/src/command/draw.rs
@@ -6,6 +6,7 @@ use wgt::error::{ErrorType, WebGpuError};
 
 use super::bind::BinderError;
 use crate::command::pass;
+use crate::validation::InvalidWorkgroupSizeError;
 use crate::{
     binding_model::{BindingError, ImmediateUploadError, LateMinBufferBindingSizeMismatch},
     resource::{
@@ -60,14 +61,8 @@ pub enum DrawError {
         if *wanted_mesh_pipeline {"standard"} else {"mesh shader"},
     )]
     WrongPipelineType { wanted_mesh_pipeline: bool },
-    #[error(
-        "Each current draw group size dimension ({current:?}) must be less or equal to {limit}, and the product must be less or equal to {max_total}"
-    )]
-    InvalidGroupSize {
-        current: [u32; 3],
-        limit: u32,
-        max_total: u32,
-    },
+    #[error(transparent)]
+    InvalidGroupSize(#[from] InvalidWorkgroupSizeError),
     #[error(
         "Mesh shader calls in multiview render passes require enabling the `EXPERIMENTAL_MESH_SHADER_MULTIVIEW` feature, and the highest bit ({highest_view_index}) in the multiview mask must be <= `Limits::max_multiview_view_count` ({max_multiviews})"
     )]
diff --git a/wgpu-core/src/command/ray_tracing.rs b/wgpu-core/src/command/ray_tracing.rs
@@ -703,9 +703,10 @@ fn iter_blas<'snatch_guard: 'buffers, 'buffers>(
                         }) {
                             input_barriers.push(barrier);
                         }
-                        if vertex_buffer.size
-                            < (mesh.size.vertex_count + mesh.first_vertex) as u64
-                                * mesh.vertex_stride
+                        if u64::from(mesh.size.vertex_count)
+                            .checked_add(u64::from(mesh.first_vertex))
+                            .and_then(|end_index| end_index.checked_mul(mesh.vertex_stride))
+                            .is_none_or(|end| vertex_buffer.size < end)
                         {
                             return Err(BuildAccelerationStructureError::InsufficientBufferSize(
                                 vertex_buffer.error_ident(),
diff --git a/wgpu-core/src/command/render.rs b/wgpu-core/src/command/render.rs
@@ -44,7 +44,8 @@ use crate::{
     },
     snatch::SnatchGuard,
     track::{ResourceUsageCompatibilityError, Tracker, UsageScope},
-    validation, Label,
+    validation::{self, check_workgroup_sizes},
+    Label,
 };
 
 #[cfg(feature = "serde")]
@@ -2731,21 +2732,17 @@ fn draw_mesh_tasks(
         .device
         .limits
         .max_task_mesh_workgroup_total_count;
-    if group_count_x > groups_size_limit
-        || group_count_y > groups_size_limit
-        || group_count_z > groups_size_limit
-        || group_count_x * group_count_y * group_count_z > max_groups
-    {
-        return Err(DrawError::InvalidGroupSize {
-            current: [group_count_x, group_count_y, group_count_z],
-            limit: groups_size_limit,
-            max_total: max_groups,
-        }
-        .into());
-    }
+    let total_count = check_workgroup_sizes(
+        &[group_count_x, group_count_y, group_count_z],
+        &[groups_size_limit, groups_size_limit, groups_size_limit],
+        "max_task_mesh_workgroups_per_dimension",
+        max_groups,
+        "max_task_mesh_workgroup_total_count",
+    )
+    .map_err(|err| RenderPassErrorInner::Draw(err.into()))?;
 
     unsafe {
-        if group_count_x > 0 && group_count_y > 0 && group_count_z > 0 {
+        if total_count > 0 {
             state.pass.base.raw_encoder.draw_mesh_tasks(
                 group_count_x,
                 group_count_y,
diff --git a/wgpu-core/src/device/queue.rs b/wgpu-core/src/device/queue.rs
@@ -868,6 +868,7 @@ impl Queue {
             self.device.alignments.buffer_copy_pitch.get() as u32,
             block_size,
         );
+        assert!(u32::MAX - bytes_in_last_row >= bytes_per_row_alignment);
         let stage_bytes_per_row = wgt::math::align_to(bytes_in_last_row, bytes_per_row_alignment);
 
         // Platform validation requires that the staging buffer always be
@@ -883,17 +884,21 @@ impl Queue {
         } else {
             profiling::scope!("copy chunked");
             // Copy row by row into the optimal alignment.
-            let block_rows_in_copy =
-                (size.depth_or_array_layers - 1) * rows_per_image + height_in_blocks;
-            let stage_size =
-                wgt::BufferSize::new(stage_bytes_per_row as u64 * block_rows_in_copy as u64)
-                    .unwrap();
+            let block_rows_in_copy = u64::from(size.depth_or_array_layers - 1)
+                * u64::from(rows_per_image)
+                + u64::from(height_in_blocks);
+            // The copy size was validated against the source buffer, however,
+            // `stage_bytes_per_row` can differ, so let's be paranoid.
+            let stage_size = u64::from(stage_bytes_per_row)
+                .checked_mul(block_rows_in_copy)
+                .and_then(wgt::BufferSize::new)
+                .unwrap();
             let mut staging_buffer = StagingBuffer::new(&self.device, stage_size)?;
-            for layer in 0..size.depth_or_array_layers {
-                let rows_offset = layer * rows_per_image;
-                for row in rows_offset..rows_offset + height_in_blocks {
-                    let src_offset = data_layout.offset as u32 + row * bytes_per_row;
-                    let dst_offset = row * stage_bytes_per_row;
+            for layer in 0..u64::from(size.depth_or_array_layers) {
+                let rows_offset = layer * u64::from(rows_per_image);
+                for row in rows_offset..rows_offset + u64::from(height_in_blocks) {
+                    let src_offset = data_layout.offset + row * u64::from(bytes_per_row);
+                    let dst_offset = row * u64::from(stage_bytes_per_row);
                     unsafe {
                         staging_buffer.write_with_offset(
                             data,
diff --git a/wgpu-core/src/validation.rs b/wgpu-core/src/validation.rs
@@ -281,18 +281,8 @@ impl WebGpuError for InputError {
 #[derive(Clone, Debug, Error)]
 #[non_exhaustive]
 pub enum StageError {
-    #[error(
-        "Shader entry point's workgroup size {current:?} ({current_total} total invocations) must be less or equal to the per-dimension
-        limit `Limits::{per_dimension_limit}` of {limit:?} and the total invocation limit `Limits::{total_limit}` of {total}"
-    )]
-    InvalidWorkgroupSize {
-        current: [u32; 3],
-        current_total: u32,
-        limit: [u32; 3],
-        total: u32,
-        per_dimension_limit: &'static str,
-        total_limit: &'static str,
-    },
+    #[error(transparent)]
+    InvalidWorkgroupSize(#[from] InvalidWorkgroupSizeError),
     #[error("Unable to find entry point '{0}'")]
     MissingEntryPoint(String),
     #[error("Shader global {0:?} is not available in the pipeline layout")]
@@ -1384,63 +1374,48 @@ impl Interface {
 
         // check workgroup size limits
         if shader_stage.to_naga().compute_like() {
-            let (
-                max_workgroup_size_limits,
-                max_workgroup_size_total,
-                per_dimension_limit,
-                total_limit,
-            ) = match shader_stage.to_naga() {
-                naga::ShaderStage::Compute => (
-                    [
+            let total = match shader_stage.to_naga() {
+                naga::ShaderStage::Compute => check_workgroup_sizes(
+                    &entry_point.workgroup_size,
+                    &[
                         self.limits.max_compute_workgroup_size_x,
                         self.limits.max_compute_workgroup_size_y,
                         self.limits.max_compute_workgroup_size_z,
                     ],
-                    self.limits.max_compute_invocations_per_workgroup,
                     "max_compute_workgroup_size_*",
+                    self.limits.max_compute_invocations_per_workgroup,
                     "max_compute_invocations_per_workgroup",
-                ),
-                naga::ShaderStage::Task => (
-                    [
+                )?,
+                naga::ShaderStage::Task => check_workgroup_sizes(
+                    &entry_point.workgroup_size,
+                    &[
                         self.limits.max_task_invocations_per_dimension,
                         self.limits.max_task_invocations_per_dimension,
                         self.limits.max_task_invocations_per_dimension,
                     ],
-                    self.limits.max_task_invocations_per_workgroup,
                     "max_task_invocations_per_dimension",
+                    self.limits.max_task_invocations_per_workgroup,
                     "max_task_invocations_per_workgroup",
-                ),
-                naga::ShaderStage::Mesh => (
-                    [
+                )?,
+                naga::ShaderStage::Mesh => check_workgroup_sizes(
+                    &entry_point.workgroup_size,
+                    &[
                         self.limits.max_mesh_invocations_per_dimension,
                         self.limits.max_mesh_invocations_per_dimension,
                         self.limits.max_mesh_invocations_per_dimension,
                     ],
-                    self.limits.max_mesh_invocations_per_workgroup,
                     "max_mesh_invocations_per_dimension",
+                    self.limits.max_mesh_invocations_per_workgroup,
                     "max_mesh_invocations_per_workgroup",
-                ),
+                )?,
                 _ => unreachable!(),
             };
-            let total_invocations = entry_point
-                .workgroup_size
-                .iter()
-                .fold(1u32, |total, &dim| total.saturating_mul(dim));
-            let invalid_total_invocations =
-                total_invocations > max_workgroup_size_total || total_invocations == 0;
-
-            let dimension_too_large = entry_point.workgroup_size[0] > max_workgroup_size_limits[0]
-                || entry_point.workgroup_size[1] > max_workgroup_size_limits[1]
-                || entry_point.workgroup_size[2] > max_workgroup_size_limits[2];
-            if invalid_total_invocations || dimension_too_large {
-                return Err(StageError::InvalidWorkgroupSize {
-                    current: entry_point.workgroup_size,
-                    current_total: total_invocations,
-                    limit: max_workgroup_size_limits,
-                    total: max_workgroup_size_total,
-                    per_dimension_limit,
-                    total_limit,
-                });
+            if total == 0 {
+                return Err(StageError::InvalidWorkgroupSize(
+                    InvalidWorkgroupSizeError::Zero {
+                        dimensions: entry_point.workgroup_size,
+                    },
+                ));
             }
         }
 
@@ -1810,6 +1785,62 @@ pub fn validate_color_attachment_bytes_per_sample(
     Ok(())
 }
 
+#[derive(Clone, Debug, Error)]
+pub enum InvalidWorkgroupSizeError {
+    #[error(
+        "Workgroup size {dimensions:?} ({total} total invocations) must be less or equal to \
+        the per-dimension limit `Limits::{per_dimension_limits_desc}` of {per_dimension_limits:?} \
+        and the total invocation limit `Limits::{total_limit_desc}` of {total_limit}"
+    )]
+    LimitExceeded {
+        dimensions: [u32; 3],
+        per_dimension_limits: [u32; 3],
+        per_dimension_limits_desc: &'static str,
+        total: u32,
+        total_limit: u32,
+        total_limit_desc: &'static str,
+    },
+    #[error("Workgroup sizes {dimensions:?} must be positive")]
+    Zero { dimensions: [u32; 3] },
+}
+
+/// Check X/Y/Z workgroup sizes against per-dimension and overall limits.
+///
+/// This function does not check that the sizes are non-zero. In a dispatch, it is legal for
+/// the size to be zero. In shader or pipeline creation, it is an error for the size to be
+/// zero, and the caller must check that.
+pub(crate) fn check_workgroup_sizes(
+    sizes: &[u32; 3],
+    per_dimension_limits: &[u32; 3],
+    per_dimension_limits_desc: &'static str,
+    total_limit: u32,
+    total_limit_desc: &'static str,
+) -> Result<u32, InvalidWorkgroupSizeError> {
+    let total = sizes
+        .iter()
+        .fold(1u32, |total, &dim| total.saturating_mul(dim));
+
+    let invalid_total_invocations = total > total_limit;
+
+    let dimension_too_large = sizes
+        .iter()
+        .zip(per_dimension_limits.iter())
+        .any(|(dim, limit)| dim > limit);
+
+    if invalid_total_invocations || dimension_too_large {
+        Err(InvalidWorkgroupSizeError::LimitExceeded {
+            dimensions: *sizes,
+            per_dimension_limits: *per_dimension_limits,
+            per_dimension_limits_desc,
+            total,
+            total_limit,
+            total_limit_desc,
+        })
+    } else {
+        Ok(total)
+    }
+}
+
 pub enum ShaderStageForValidation {
     Vertex {
         topology: wgt::PrimitiveTopology,

Original file line number	Diff line number	Diff line change
`@@ -885,11 +885,11 @@ fn dispatch_indirect(`
`885`	`885`	`return Err(ComputePassErrorInner::UnalignedIndirectBufferOffset(offset));`
`886`	`886`	`}`
`887`	`887`
`888`		`- let end_offset = offset + size_of::<wgt::DispatchIndirectArgs>() as u64;`
`889`		`- if end_offset > buffer.size {`
	`888`	`+ let args_size = size_of::<wgt::DispatchIndirectArgs>() as u64;`
	`889`	`+ if buffer.size < args_size \|\| buffer.size - args_size < offset {`
`890`	`890`	`return Err(ComputePassErrorInner::IndirectBufferOverrun {`
`891`	`891`	`offset,`
`892`		`- end_offset,`
	`892`	`+ end_offset: offset + args_size,`
`893`	`893`	`buffer_size: buffer.size,`
`894`	`894`	`});`
`895`	`895`	`}`
Original file line number	Diff line number	Diff line change
`@@ -703,9 +703,10 @@ fn iter_blas<'snatch_guard: 'buffers, 'buffers>(`
`703`	`703`	`}) {`
`704`	`704`	`input_barriers.push(barrier);`
`705`	`705`	`}`
`706`		`- if vertex_buffer.size`
`707`		`- < (mesh.size.vertex_count + mesh.first_vertex) as u64`
`708`		`- * mesh.vertex_stride`
	`706`	`+ if u64::from(mesh.size.vertex_count)`
	`707`	`+ .checked_add(u64::from(mesh.first_vertex))`
	`708`	`+ .and_then(\|end_index\| end_index.checked_mul(mesh.vertex_stride))`
	`709`	`+ .is_none_or(\|end\| vertex_buffer.size < end)`
`709`	`710`	`{`
`710`	`711`	`return Err(BuildAccelerationStructureError::InsufficientBufferSize(`
`711`	`712`	`vertex_buffer.error_ident(),`