Refactor check_workgroup_sizes into a new struct WorkgroupSizeCheck

wgpu-core/src/command/bundle.rs

@@ -121,8 +121,8 @@ use crate::{
     snatch::SnatchGuard,
     track::RenderBundleScope,
     validation::{
-        check_color_attachment_count, check_workgroup_sizes,
-        validate_color_attachment_bytes_per_sample,
+        check_color_attachment_count, validate_color_attachment_bytes_per_sample,
+        WorkgroupSizeCheck,
     },
     Label, LabelHelpers,
 };
@@ -912,13 +912,15 @@ fn draw_mesh_tasks(
         )
     };
 
-    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",
-    )
+    let total_count = WorkgroupSizeCheck {
+        dimensions: &[group_count_x, group_count_y, group_count_z],
+        per_dimension_limits: &[groups_size_limit, groups_size_limit, groups_size_limit],
+        per_dimension_limits_desc: "max_task_mesh_workgroups_per_dimension",
+
+        total_limit: max_groups,
+        total_limit_desc: "max_task_mesh_workgroup_total_count",
+    }
+    .check_and_compute_total_invocations()
     .map_err(|err| RenderBundleErrorInner::Draw(err.into()))?;
 
     if total_count > 0 {

wgpu-core/src/command/render.rs

@@ -45,7 +45,7 @@ use crate::{
     },
     snatch::SnatchGuard,
     track::{ResourceUsageCompatibilityError, Tracker, UsageScope},
-    validation::{self, check_workgroup_sizes},
+    validation::{self, WorkgroupSizeCheck},
     Label,
 };
 
@@ -3099,13 +3099,15 @@ fn draw_mesh_tasks(
         )
     };
 
-    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",
-    )
+    let total_count = WorkgroupSizeCheck {
+        dimensions: &[group_count_x, group_count_y, group_count_z],
+        per_dimension_limits: &[groups_size_limit, groups_size_limit, groups_size_limit],
+        per_dimension_limits_desc: "max_task_mesh_workgroups_per_dimension",
+
+        total_limit: max_groups,
+        total_limit_desc: "max_task_mesh_workgroup_total_count",
+    }
+    .check_and_compute_total_invocations()
     .map_err(|err| RenderPassErrorInner::Draw(err.into()))?;
 
     unsafe {

wgpu-core/src/validation.rs

@@ -1505,42 +1505,38 @@ impl Interface {
 
         // check workgroup size limits
         if shader_stage.to_naga().compute_like() {
-            let total = match shader_stage.to_naga() {
-                naga::ShaderStage::Compute => check_workgroup_sizes(
-                    &entry_point.workgroup_size,
-                    &[
+            let workgroup_size_check = match shader_stage.to_naga() {
+                naga::ShaderStage::Compute => WorkgroupSizeCheck {
+                    dimensions: &entry_point.workgroup_size,
+                    per_dimension_limits: &[
                         self.limits.max_compute_workgroup_size_x,
                         self.limits.max_compute_workgroup_size_y,
                         self.limits.max_compute_workgroup_size_z,
                     ],
-                    "max_compute_workgroup_size_*",
-                    self.limits.max_compute_invocations_per_workgroup,
-                    "max_compute_invocations_per_workgroup",
-                )?,
-                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,
-                    ],
-                    "max_task_invocations_per_dimension",
-                    self.limits.max_task_invocations_per_workgroup,
-                    "max_task_invocations_per_workgroup",
-                )?,
-                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,
-                    ],
-                    "max_mesh_invocations_per_dimension",
-                    self.limits.max_mesh_invocations_per_workgroup,
-                    "max_mesh_invocations_per_workgroup",
-                )?,
+                    per_dimension_limits_desc: "max_compute_workgroup_size_*",
+
+                    total_limit: self.limits.max_compute_invocations_per_workgroup,
+                    total_limit_desc: "max_compute_invocations_per_workgroup",
+                },
+                naga::ShaderStage::Task => WorkgroupSizeCheck {
+                    dimensions: &entry_point.workgroup_size,
+                    per_dimension_limits: &[self.limits.max_task_invocations_per_dimension; 3],
+                    per_dimension_limits_desc: "max_task_invocations_per_dimension",
+
+                    total_limit: self.limits.max_task_invocations_per_workgroup,
+                    total_limit_desc: "max_task_invocations_per_workgroup",
+                },
+                naga::ShaderStage::Mesh => WorkgroupSizeCheck {
+                    dimensions: &entry_point.workgroup_size,
+                    per_dimension_limits: &[self.limits.max_mesh_invocations_per_dimension; 3],
+                    per_dimension_limits_desc: "max_mesh_invocations_per_dimension",
+
+                    total_limit: self.limits.max_mesh_invocations_per_workgroup,
+                    total_limit_desc: "max_mesh_invocations_per_workgroup",
+                },
                 _ => unreachable!(),
             };
+            let total = workgroup_size_check.check_and_compute_total_invocations()?;
             if total == 0 {
                 return Err(StageError::InvalidWorkgroupSize(
                     InvalidWorkgroupSizeError::Zero {
@@ -1968,40 +1964,57 @@ pub enum InvalidWorkgroupSizeError {
     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,
+/// A helper type for avoiding argument order mistakes when calling
+/// [`Self::check_and_compute_total_invocations`].
+#[derive(Clone, Debug)]
+pub(crate) struct WorkgroupSizeCheck<'a> {
+    pub dimensions: &'a [u32; 3],
+    pub per_dimension_limits: &'a [u32; 3],
+    pub per_dimension_limits_desc: &'static str,
+    pub total_limit: u32,
+    pub total_limit_desc: &'static str,
+}
+
+impl WorkgroupSizeCheck<'_> {
+    /// 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_and_compute_total_invocations(
+        self,
+    ) -> Result<u32, InvalidWorkgroupSizeError> {
+        let Self {
+            dimensions,
+            per_dimension_limits,
             per_dimension_limits_desc,
-            total,
             total_limit,
             total_limit_desc,
-        })
-    } else {
-        Ok(total)
+        } = self;
+
+        let total = dimensions
+            .iter()
+            .fold(1u32, |total, &dim| total.saturating_mul(dim));
+
+        let invalid_total_invocations = total > total_limit;
+
+        let dimension_too_large = dimensions
+            .iter()
+            .zip(per_dimension_limits.iter())
+            .any(|(dim, limit)| dim > limit);
+
+        if invalid_total_invocations || dimension_too_large {
+            Err(InvalidWorkgroupSizeError::LimitExceeded {
+                dimensions: *dimensions,
+                per_dimension_limits: *per_dimension_limits,
+                per_dimension_limits_desc,
+                total,
+                total_limit,
+                total_limit_desc,
+            })
+        } else {
+            Ok(total)
+        }
     }
 }