From ba5292a2c05fac958087b6627e1c3d2a7d8a8b92 Mon Sep 17 00:00:00 2001 From: Jim Blandy Date: Tue, 7 Nov 2023 13:06:53 -0800 Subject: [PATCH] [wgsl-in] Remove unused expressions from Module::const_expressions. After building a Naga module from the WGSL AST, make a pass over the module to determine which expressions in `Module::const_expressions` are actually used, and compact out all unused expressions from that arena, adjusting references appropriately. This isn't necessary on trunk, but it will be when we add abstract types to the WGSL front end. The evaluation of expressions containing abstract values will naturally introduce lots of 64-bit float and integer literals to the constant arena, which the validator would rightly reject. However, because all abstract values are converted to concrete types before appearing as part of the module's final code, all such 64-bit expressions should be unused. --- CHANGELOG.md | 5 + naga/src/arena.rs | 2 +- naga/src/front/wgsl/lower/compact.rs | 160 +++++++++++++++++++++++++++ naga/src/front/wgsl/lower/mod.rs | 3 + naga/tests/out/ir/access.ron | 26 ----- naga/tests/out/ir/collatz.ron | 6 +- 6 files changed, 170 insertions(+), 32 deletions(-) create mode 100644 naga/src/front/wgsl/lower/compact.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 7b1190c904..46d79ec0d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -58,6 +58,11 @@ For naga changelogs at or before v0.14.0. See [naga's changelog](naga/CHANGELOG. - Fix issue where local variables were sometimes using variable names from previous functions. +### Changes +#### Naga + +- [wgsl-in] Omit unused constant expressions from the Module's `const_expressions` arena, in preparation for supporting abstract types. By @jimblandy in [#4648](https://github.com/gfx-rs/wgpu/pull/4648). + ## v0.18.0 (2023-10-25) ### Desktop OpenGL 3.3+ Support on Windows diff --git a/naga/src/arena.rs b/naga/src/arena.rs index a9df066654..009d65b1c3 100644 --- a/naga/src/arena.rs +++ b/naga/src/arena.rs @@ -424,7 +424,7 @@ impl Arena { Ok(()) } - #[cfg(feature = "compact")] + #[cfg(any(feature = "compact", feature = "wgsl-in"))] pub(crate) fn retain_mut

(&mut self, mut predicate: P) where P: FnMut(Handle, &mut T) -> bool, diff --git a/naga/src/front/wgsl/lower/compact.rs b/naga/src/front/wgsl/lower/compact.rs new file mode 100644 index 0000000000..9339411f2b --- /dev/null +++ b/naga/src/front/wgsl/lower/compact.rs @@ -0,0 +1,160 @@ +use crate::Handle; +use bit_set::BitSet; + +/// Remove all unused expressions from `module.const_expressions`. +/// +/// Abstract values make extensive use of 64-bit literals, which the +/// validator should never see. +pub fn compact(module: &mut crate::Module) { + // Trace which expressions in `module.const_expressions` are actually used. + let used = trace(module); + + // Assuming unused expressions are squeezed out of the arena, + // compute a map from pre-squeeze indices to post-squeeze + // expression handles. + let mut next = 0; + let compacted: Vec>> = module + .const_expressions + .iter() + .map(|(handle, _)| { + if used.contains(handle.index()) { + let index = next; + next += 1; + std::num::NonZeroU32::new(index + 1).map(Handle::new) + } else { + None + } + }) + .collect(); + + adjust(module, &compacted); +} + +fn trace(module: &crate::Module) -> BitSet { + let mut used = BitSet::new(); + + // Note uses by global constants. + for (_, constant) in module.constants.iter() { + used.insert(constant.init.index()); + } + + // Note uses by global variable initializers. + for (_, global) in module.global_variables.iter() { + if let Some(init) = global.init { + used.insert(init.index()); + } + } + + // Note uses by functions' expressions. + for (_, fun) in module.functions.iter() { + trace_function(fun, &mut used); + } + + // Note uses by entry points' expressions. + for entry_point in &module.entry_points { + trace_function(&entry_point.function, &mut used); + } + + // Note transitive uses by other used constant expressions. + // + // Since we know that expressions only refer to other expressions + // appearing before them in the arena, we can do this without a + // stack by making a single pass from back to front. + for (handle, expr) in module.const_expressions.iter().rev() { + if used.contains(handle.index()) { + match *expr { + crate::Expression::Compose { ref components, .. } => { + for component in components { + used.insert(component.index()); + } + } + crate::Expression::Splat { value, .. } => { + used.insert(value.index()); + } + _ => {} + } + } + } + + used +} + +fn trace_function(function: &crate::Function, used: &mut BitSet) { + for (_, expr) in function.expressions.iter() { + match *expr { + crate::Expression::ImageSample { + offset: Some(offset), + .. + } => { + used.insert(offset.index()); + } + _ => {} + } + } +} + +fn adjust(module: &mut crate::Module, compacted: &[Option>]) { + // Remove unused expressions from the constant arena, + // and adjust the handles in retained expressions. + module.const_expressions.retain_mut(|handle, expr| { + match compacted[handle.index()] { + Some(_) => { + // This expression is used, and thus its handles are worth adjusting. + match *expr { + crate::Expression::Compose { + ref mut components, .. + } => { + for component in components { + *component = compacted[component.index()].unwrap(); + } + } + crate::Expression::Splat { ref mut value, .. } => { + *value = compacted[value.index()].unwrap(); + } + _ => {} + } + true + } + None => false, + } + }); + + // Adjust uses by global constants. + for (_, constant) in module.constants.iter_mut() { + constant.init = compacted[constant.init.index()].unwrap(); + } + + // Adjust uses by global variable initializers. + for (_, global) in module.global_variables.iter_mut() { + if let Some(ref mut init) = global.init { + *init = compacted[init.index()].unwrap(); + } + } + + // Adjust uses by functions' expressions. + for (_, fun) in module.functions.iter_mut() { + adjust_function(fun, compacted); + } + + // Adjust uses by entry points' expressions. + for entry_point in &mut module.entry_points { + adjust_function(&mut entry_point.function, compacted); + } +} + +fn adjust_function( + function: &mut crate::Function, + compacted: &[Option>], +) { + for (_, expr) in function.expressions.iter_mut() { + match *expr { + crate::Expression::ImageSample { + offset: Some(ref mut offset), + .. + } => { + *offset = compacted[offset.index()].unwrap(); + } + _ => {} + } + } +} diff --git a/naga/src/front/wgsl/lower/mod.rs b/naga/src/front/wgsl/lower/mod.rs index 567360b5d8..b1294eb438 100644 --- a/naga/src/front/wgsl/lower/mod.rs +++ b/naga/src/front/wgsl/lower/mod.rs @@ -11,6 +11,7 @@ use crate::proc::{ }; use crate::{Arena, FastHashMap, FastIndexMap, Handle, Span}; +mod compact; mod construction; /// Resolves the inner type of a given expression. @@ -971,6 +972,8 @@ impl<'source, 'temp> Lowerer<'source, 'temp> { } } + compact::compact(&mut module); + Ok(module) } diff --git a/naga/tests/out/ir/access.ron b/naga/tests/out/ir/access.ron index 77d95dd58f..16195df90c 100644 --- a/naga/tests/out/ir/access.ron +++ b/naga/tests/out/ir/access.ron @@ -412,32 +412,6 @@ 6, ], ), - Literal(I32(8)), - Literal(I32(2)), - Literal(I32(10)), - Literal(I32(2)), - Literal(I32(0)), - Literal(I32(0)), - Literal(I32(0)), - Literal(I32(1)), - Literal(I32(0)), - Literal(I32(2)), - Literal(I32(2)), - Literal(I32(0)), - Literal(I32(3)), - Literal(I32(2)), - Literal(I32(2)), - Literal(I32(10)), - Literal(I32(5)), - Literal(I32(5)), - Literal(I32(10)), - Literal(I32(5)), - Literal(I32(0)), - Literal(I32(2)), - Literal(I32(2)), - Literal(I32(2)), - Literal(I32(2)), - Literal(I32(1)), ], functions: [ ( diff --git a/naga/tests/out/ir/collatz.ron b/naga/tests/out/ir/collatz.ron index 8146909c1e..7cad54b713 100644 --- a/naga/tests/out/ir/collatz.ron +++ b/naga/tests/out/ir/collatz.ron @@ -58,11 +58,7 @@ init: None, ), ], - const_expressions: [ - Literal(I32(0)), - Literal(I32(0)), - Literal(I32(1)), - ], + const_expressions: [], functions: [ ( name: Some("collatz_iterations"),