Add Pipeline Overrides for workgroup_size

CHANGELOG.md

@@ -102,6 +102,7 @@ By @ErichDonGubler in [#6456](https://github.com/gfx-rs/wgpu/pull/6456), [#6148]
 - Implement `quantizeToF16()` for WGSL frontend, and WGSL, SPIR-V, HLSL, MSL, and GLSL backends. By @jamienicol in [#6519](https://github.com/gfx-rs/wgpu/pull/6519).
 - Add support for GLSL `usampler*` and `isampler*`. By @DavidPeicho in [#6513](https://github.com/gfx-rs/wgpu/pull/6513).
 - Expose Ray Query flags as constants in WGSL. Implement candidate intersections. By @kvark in [#5429](https://github.com/gfx-rs/wgpu/pull/5429)
+- Allow for override-expressions in `workgroup_size`. By @KentSlaney in [#6635](https://github.com/gfx-rs/wgpu/pull/6635).
 
 #### General
 

naga/src/back/pipeline_constants.rs

@@ -25,6 +25,8 @@ pub enum PipelineConstantError {
     ConstantEvaluatorError(#[from] ConstantEvaluatorError),
     #[error(transparent)]
     ValidationError(#[from] WithSpan<ValidationError>),
+    #[error("workgroup_size was overridden to a negative value")]
+    NegativeWorkgroupSize,
 }
 
 /// Replace all overrides in `module` with constants.
@@ -190,6 +192,7 @@ pub fn process_overrides<'a>(
     let mut entry_points = mem::take(&mut module.entry_points);
     for ep in entry_points.iter_mut() {
         process_function(&mut module, &override_map, &mut ep.function)?;
+        process_workgroup_size_override(&mut module, &override_map, ep)?;
     }
     module.entry_points = entry_points;
 
@@ -202,6 +205,52 @@ pub fn process_overrides<'a>(
     Ok((Cow::Owned(module), Cow::Owned(module_info)))
 }
 
+fn process_workgroup_size_override(
+    module: &mut Module,
+    override_map: &HandleVec<Override, Handle<Constant>>,
+    ep: &mut crate::EntryPoint,
+) -> Result<(), PipelineConstantError> {
+    match ep.workgroup_size_overrides {
+        None => {}
+        Some(overrides) => {
+            overrides.iter().enumerate().try_for_each(
+                |(i, overridden)| -> Result<(), PipelineConstantError> {
+                    match *overridden {
+                        None => Ok(()),
+                        Some(h) => {
+                            let c = module.constants[override_map[h]].init;
+                            let n = &module.global_expressions[c];
+                            match *n {
+                                crate::Expression::Literal(literal) => {
+                                    ep.workgroup_size[i] = match literal {
+                                        crate::Literal::U32(m) => m,
+                                        crate::Literal::I32(m) => {
+                                            if m < 0 {
+                                                Err(PipelineConstantError::NegativeWorkgroupSize)?;
+                                                unreachable!();
+                                            } else {
+                                                m as u32
+                                            }
+                                        }
+                                        _ => {
+                                            unreachable!();
+                                        }
+                                    };
+                                }
+                                _ => {
+                                    unreachable!();
+                                }
+                            }
+                            Ok(())
+                        }
+                    }
+                },
+            )?;
+        }
+    }
+    Ok(())
+}
+
 /// Add a [`Constant`] to `module` for the override `old_h`.
 ///
 /// Add the new `Constant` to `override_map` and `adjusted_constant_initializers`.

naga/src/front/glsl/functions.rs

@@ -1366,6 +1366,7 @@ impl Frontend {
             early_depth_test: Some(crate::EarlyDepthTest { conservative: None })
                 .filter(|_| self.meta.early_fragment_tests),
             workgroup_size: self.meta.workgroup_size,
+            workgroup_size_overrides: None,
             function: Function {
                 arguments,
                 expressions,

naga/src/front/spv/function.rs

@@ -569,6 +569,7 @@ impl<I: Iterator<Item = u32>> super::Frontend<I> {
             stage: ep.stage,
             early_depth_test: ep.early_depth_test,
             workgroup_size: ep.workgroup_size,
+            workgroup_size_overrides: None,
             function,
         });
 

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

@@ -1311,24 +1311,54 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             .collect();
 
         if let Some(ref entry) = f.entry_point {
-            let workgroup_size = if let Some(workgroup_size) = entry.workgroup_size {
+            let workgroup_size_info = if let Some(workgroup_size) = entry.workgroup_size {
                 // TODO: replace with try_map once stabilized
                 let mut workgroup_size_out = [1; 3];
+                let mut workgroup_size_overrides_out = [None; 3];
                 for (i, size) in workgroup_size.into_iter().enumerate() {
                     if let Some(size_expr) = size {
-                        workgroup_size_out[i] = self.const_u32(size_expr, &mut ctx.as_const())?.0;
+                        match self.const_u32(size_expr, &mut ctx.as_const()) {
+                            Ok(value) => {
+                                workgroup_size_out[i] = value.0;
+                            }
+                            err => {
+                                if let Err(Error::ConstantEvaluatorError(ref ty, _)) = err {
+                                    match **ty {
+                                        crate::proc::ConstantEvaluatorError::OverrideExpr => {
+                                            workgroup_size_overrides_out[i] =
+                                                Some(self.workgroup_size_override(
+                                                    size_expr,
+                                                    &mut ctx.as_override(),
+                                                    i,
+                                                )?);
+                                        }
+                                        _ => {
+                                            err?;
+                                        }
+                                    }
+                                } else {
+                                    err?;
+                                }
+                            }
+                        }
                     }
                 }
-                workgroup_size_out
+                if workgroup_size_overrides_out.iter().all(|x| x.is_none()) {
+                    (workgroup_size_out, None)
+                } else {
+                    (workgroup_size_out, Some(workgroup_size_overrides_out))
+                }
             } else {
-                [0; 3]
+                ([0; 3], None)
             };
 
+            let (workgroup_size, workgroup_size_overrides) = workgroup_size_info;
             ctx.module.entry_points.push(crate::EntryPoint {
                 name: f.name.name.to_string(),
                 stage: entry.stage,
                 early_depth_test: entry.early_depth_test,
                 workgroup_size,
+                workgroup_size_overrides,
                 function,
             });
             Ok(LoweredGlobalDecl::EntryPoint)
@@ -1338,6 +1368,31 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         }
     }
 
+    fn workgroup_size_override(
+        &mut self,
+        size_expr: Handle<ast::Expression<'source>>,
+        ctx: &mut ExpressionContext<'source, '_, '_>,
+        i: usize,
+    ) -> Result<Handle<crate::Override>, Error<'source>> {
+        let span = ctx.ast_expressions.get_span(size_expr);
+        let expr = self.expression(size_expr, ctx)?;
+        let ty = ctx.register_type(expr)?;
+        match ctx.module.types[ty].inner.scalar_kind().ok_or(0) {
+            Ok(crate::ScalarKind::Sint) | Ok(crate::ScalarKind::Uint) => Ok({
+                ctx.module.overrides.append(
+                    crate::Override {
+                        name: Some(format!("__workgroup_size_{}", i)),
+                        id: None,
+                        ty,
+                        init: Some(expr),
+                    },
+                    span,
+                )
+            }),
+            _ => Err(Error::ExpectedConstExprConcreteIntegerScalar(span)),
+        }
+    }
+
     fn block(
         &mut self,
         b: &ast::Block<'source>,

naga/src/lib.rs

@@ -2186,6 +2186,8 @@ pub struct EntryPoint {
     pub early_depth_test: Option<EarlyDepthTest>,
     /// Workgroup size for compute stages
     pub workgroup_size: [u32; 3],
+    /// Override expressions for workgroup size
+    pub workgroup_size_overrides: Option<[Option<Handle<Override>>; 3]>,
     /// The entrance function.
     pub function: Function,
 }

naga/tests/out/ir/access.compact.ron

@@ -1854,6 +1854,7 @@
             stage: Vertex,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("foo_vert"),
                 arguments: [
@@ -2156,6 +2157,7 @@
             stage: Fragment,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("foo_frag"),
                 arguments: [],
@@ -2348,6 +2350,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("assign_through_ptr"),
                 arguments: [],
@@ -2430,6 +2433,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("assign_to_ptr_components"),
                 arguments: [],

naga/tests/out/ir/access.ron

@@ -1854,6 +1854,7 @@
             stage: Vertex,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("foo_vert"),
                 arguments: [
@@ -2156,6 +2157,7 @@
             stage: Fragment,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("foo_frag"),
                 arguments: [],
@@ -2348,6 +2350,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("assign_through_ptr"),
                 arguments: [],
@@ -2430,6 +2433,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("assign_to_ptr_components"),
                 arguments: [],

naga/tests/out/ir/atomic_i_increment.compact.ron

@@ -263,6 +263,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (32, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("stage::test_atomic_i_increment_wrap"),
                 arguments: [],

naga/tests/out/ir/atomic_i_increment.ron

@@ -288,6 +288,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (32, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("stage::test_atomic_i_increment_wrap"),
                 arguments: [],

naga/tests/out/ir/collatz.compact.ron

@@ -257,6 +257,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("main"),
                 arguments: [

naga/tests/out/ir/collatz.ron

@@ -257,6 +257,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("main"),
                 arguments: [

naga/tests/out/ir/fetch_depth.compact.ron

@@ -176,6 +176,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (32, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("cull::fetch_depth_wrap"),
                 arguments: [],

naga/tests/out/ir/fetch_depth.ron

@@ -246,6 +246,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (32, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("cull::fetch_depth_wrap"),
                 arguments: [],

naga/tests/out/ir/index-by-value.compact.ron

@@ -300,6 +300,7 @@
             stage: Vertex,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("index_let_array_1d"),
                 arguments: [

naga/tests/out/ir/index-by-value.ron

@@ -300,6 +300,7 @@
             stage: Vertex,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("index_let_array_1d"),
                 arguments: [

naga/tests/out/ir/overrides-atomicCompareExchangeWeak.compact.ron

@@ -85,6 +85,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("f"),
                 arguments: [],

naga/tests/out/ir/overrides-atomicCompareExchangeWeak.ron

@@ -85,6 +85,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("f"),
                 arguments: [],

naga/tests/out/ir/overrides-ray-query.compact.ron

@@ -111,6 +111,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("main"),
                 arguments: [],

naga/tests/out/ir/overrides-ray-query.ron

@@ -111,6 +111,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("main"),
                 arguments: [],

naga/tests/out/ir/overrides.compact.ron

@@ -108,6 +108,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("main"),
                 arguments: [],

naga/tests/out/ir/overrides.ron

@@ -108,6 +108,7 @@
             stage: Compute,
             early_depth_test: None,
             workgroup_size: (1, 1, 1),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("main"),
                 arguments: [],

naga/tests/out/ir/shadow.compact.ron

@@ -958,6 +958,7 @@
             stage: Fragment,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("fs_main_wrap"),
                 arguments: [

naga/tests/out/ir/shadow.ron

@@ -1236,6 +1236,7 @@
             stage: Fragment,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("fs_main_wrap"),
                 arguments: [

naga/tests/out/ir/spec-constants.compact.ron

@@ -495,6 +495,7 @@
             stage: Vertex,
             early_depth_test: None,
             workgroup_size: (0, 0, 0),
+            workgroup_size_overrides: None,
             function: (
                 name: Some("main_wrap"),
                 arguments: [