[naga] Remove non-essential override references via compaction

.github/workflows/ci.yml

@@ -304,7 +304,7 @@ jobs:
 
           # Check with all compatible features
           cargo clippy --target ${{ matrix.target }} ${{ matrix.extra-flags }} -p wgpu-types --no-default-features --features strict_asserts,fragile-send-sync-non-atomic-wasm,serde,counters
-          cargo clippy --target ${{ matrix.target }} ${{ matrix.extra-flags }} -p naga --no-default-features --features dot-out,compact
+          cargo clippy --target ${{ matrix.target }} ${{ matrix.extra-flags }} -p naga --no-default-features --features dot-out
           cargo clippy --target ${{ matrix.target }} ${{ matrix.extra-flags }} -p wgpu-hal --no-default-features --features fragile-send-sync-non-atomic-wasm
           cargo clippy --target ${{ matrix.target }} ${{ matrix.extra-flags }} -p wgpu --no-default-features --features serde
 

CHANGELOG.md

@@ -67,8 +67,7 @@ Bottom level categories:
 
 #### Naga
 
-Naga now infers the correct binding layout when a resource appears only in an assignment to `_`. By @andyleiserson in [#7540](https://github.com/gfx-rs/wgpu/pull/7540).
-
+- Naga now infers the correct binding layout when a resource appears only in an assignment to `_`. By @andyleiserson in [#7540](https://github.com/gfx-rs/wgpu/pull/7540).
 - Implement `dot4U8Packed` and `dot4I8Packed` for all backends, using specialized intrinsics on SPIR-V, HSLS, and Metal if available, and polyfills everywhere else. By @robamler in [#7494](https://github.com/gfx-rs/wgpu/pull/7494), [#7574](https://github.com/gfx-rs/wgpu/pull/7574), and [#7653](https://github.com/gfx-rs/wgpu/pull/7653).
 - Add polyfilled `pack4x{I,U}8Clamped` built-ins to all backends and WGSL frontend. By @ErichDonGubler in [#7546](https://github.com/gfx-rs/wgpu/pull/7546).
 - Allow textureLoad's sample index arg to be unsigned. By @jimblandy in [#7625](https://github.com/gfx-rs/wgpu/pull/7625).
@@ -78,6 +77,7 @@ Naga now infers the correct binding layout when a resource appears only in an as
 - Properly evaluate `abs(most negative abstract int)`. By @jimblandy in [#7507](https://github.com/gfx-rs/wgpu/pull/7507).
 - Generate vectorized code for `[un]pack4x{I,U}8[Clamp]` on SPIR-V and MSL 2.1+. By @robamler in [#7664](https://github.com/gfx-rs/wgpu/pull/7664).
 - Fix typing for `select`, which had issues particularly with a lack of automatic type conversion. By @ErichDonGubler in [#7572](https://github.com/gfx-rs/wgpu/pull/7572).
+- Allow scalars as the first argument of the `distance` built-in function. By @bernhl in [#7530](https://github.com/gfx-rs/wgpu/pull/7530).
 
 #### DX12
 
@@ -102,6 +102,10 @@ Naga now infers the correct binding layout when a resource appears only in an as
 - `naga::back::hlsl::Writer::new` has a new `pipeline_options` argument. `hlsl::PipelineOptions::default()` can be passed as a default. The `shader_stage` and `entry_point` members of `pipeline_options` can be used to write only a single entry point when using the HLSL and MSL backends (GLSL and SPIR-V already had this functionality). The Metal and DX12 HALs now write only a single entry point when loading shaders. By @andyleiserson in [#7626](https://github.com/gfx-rs/wgpu/pull/7626).
 - Implemented `early_depth_test` for SPIR-V backend, enabling `SHADER_EARLY_DEPTH_TEST` for Vulkan. Additionally, fixed conservative depth optimizations when using `early_depth_test`. The syntax for forcing early depth tests is now `@early_depth_test(force)` instead of `@early_depth_test`. By @dzamkov in [#7676](https://github.com/gfx-rs/wgpu/pull/7676).
 - `ImplementedLanguageExtension::VARIANTS` is now implemented manually rather than derived using `strum` (allowing `strum` to become a dev-only dependency) so it is no longer a member of the `strum::VARIANTS` trait. Unless you are using this trait as a bound this should have no effect.
+- Compaction changes, by @andyleiserson in [#7703](https://github.com/gfx-rs/wgpu/pull/7703):
+  - [`process_overrides`](https://docs.rs/naga/latest/naga/back/pipeline_constants/fn.process_overrides.html) now compacts the module to remove unused items. It is no longer necessary to supply values for overrides that are not used by the active entry point.
+  - The `compact` Cargo feature has been removed. It is no longer possible to exclude compaction support from the build.
+  - [`compact`](https://docs.rs/naga/latest/naga/compact/fn.compact.html) now has an additional argument that specifies whether to remove unused functions, globals, and named types and overrides. For the previous behavior, pass `KeepUnused::Yes`.
 
 #### D3D12
 
@@ -119,11 +123,6 @@ Naga now infers the correct binding layout when a resource appears only in an as
 
 - Added initial `no_std` support to `wgpu-hal`. By @bushrat011899 in [#7599](https://github.com/gfx-rs/wgpu/pull/7599)
 
-### Bug Fixes
-
-#### Naga
-
-- Allow scalars as the first argument of the `distance` built-in function. By @bernhl in [#7530](https://github.com/gfx-rs/wgpu/pull/7530).
 
 ### Documentation
 

benches/benches/wgpu-benchmark/shader.rs

@@ -269,6 +269,8 @@ fn validation(c: &mut Criterion) {
 }
 
 fn compact(c: &mut Criterion) {
+    use naga::compact::{compact, KeepUnused};
+
     let mut inputs = get_wgsl_inputs();
 
     inputs.validate();
@@ -279,7 +281,7 @@ fn compact(c: &mut Criterion) {
     group.bench_function("shader: compact", |b| {
         b.iter(|| {
             for input in &mut inputs.inner {
-                naga::compact::compact(input.module.as_mut().unwrap());
+                compact(input.module.as_mut().unwrap(), KeepUnused::No);
             }
         });
     });

cts_runner/test.lst

@@ -7,7 +7,9 @@ webgpu:api,operation,command_buffer,copyBufferToBuffer:state_transitions:*
 webgpu:api,operation,command_buffer,copyBufferToBuffer:copy_order:*
 webgpu:api,operation,compute,basic:memcpy:*
 //FAIL: webgpu:api,operation,compute,basic:large_dispatch:*
+webgpu:api,operation,compute_pipeline,overrides:*
 webgpu:api,operation,device,lost:*
+webgpu:api,operation,render_pipeline,overrides:*
 webgpu:api,operation,rendering,basic:clear:*
 webgpu:api,operation,rendering,basic:fullscreen_quad:*
 //FAIL: webgpu:api,operation,rendering,basic:large_draw:*

naga-cli/Cargo.toml

@@ -25,7 +25,6 @@ test = false
 
 [dependencies]
 naga = { workspace = true, features = [
-    "compact",
     "wgsl-in",
     "wgsl-out",
     "glsl-in",

naga-cli/src/bin/naga.rs

@@ -542,7 +542,7 @@ fn run() -> anyhow::Result<()> {
                 write_output(&module, &info, &params, before_compaction)?;
             }
 
-            naga::compact::compact(&mut module);
+            naga::compact::compact(&mut module, KeepUnused::No);
 
             // Re-validate the IR after compaction.
             match naga::valid::Validator::new(params.validation_flags, validation_caps)
@@ -717,9 +717,13 @@ fn write_output(
                  succeed, and it failed in a previous step",
             ))?;
 
-            let (module, info) =
-                naga::back::pipeline_constants::process_overrides(module, info, &params.overrides)
-                    .unwrap_pretty();
+            let (module, info) = naga::back::pipeline_constants::process_overrides(
+                module,
+                info,
+                None,
+                &params.overrides,
+            )
+            .unwrap_pretty();
 
             let pipeline_options = msl::PipelineOptions::default();
             let (msl, _) =
@@ -751,9 +755,13 @@ fn write_output(
                  succeed, and it failed in a previous step",
             ))?;
 
-            let (module, info) =
-                naga::back::pipeline_constants::process_overrides(module, info, &params.overrides)
-                    .unwrap_pretty();
+            let (module, info) = naga::back::pipeline_constants::process_overrides(
+                module,
+                info,
+                None,
+                &params.overrides,
+            )
+            .unwrap_pretty();
 
             let spv =
                 spv::write_vec(&module, &info, &params.spv_out, pipeline_options).unwrap_pretty();
@@ -788,9 +796,13 @@ fn write_output(
                  succeed, and it failed in a previous step",
             ))?;
 
-            let (module, info) =
-                naga::back::pipeline_constants::process_overrides(module, info, &params.overrides)
-                    .unwrap_pretty();
+            let (module, info) = naga::back::pipeline_constants::process_overrides(
+                module,
+                info,
+                None,
+                &params.overrides,
+            )
+            .unwrap_pretty();
 
             let mut buffer = String::new();
             let mut writer = glsl::Writer::new(
@@ -819,9 +831,13 @@ fn write_output(
                  succeed, and it failed in a previous step",
             ))?;
 
-            let (module, info) =
-                naga::back::pipeline_constants::process_overrides(module, info, &params.overrides)
-                    .unwrap_pretty();
+            let (module, info) = naga::back::pipeline_constants::process_overrides(
+                module,
+                info,
+                None,
+                &params.overrides,
+            )
+            .unwrap_pretty();
 
             let mut buffer = String::new();
             let pipeline_options = Default::default();
@@ -906,4 +922,4 @@ fn bulk_validate(args: Args, params: &Parameters) -> anyhow::Result<()> {
 }
 
 use codespan_reporting::term::termcolor::{ColorChoice, StandardStream};
-use naga::FastHashMap;
+use naga::{compact::KeepUnused, FastHashMap};

naga/Cargo.toml

@@ -58,7 +58,7 @@ arbitrary = [
 ]
 spv-in = ["dep:petgraph", "petgraph/graphmap", "dep:spirv"]
 spv-out = ["dep:spirv"]
-wgsl-in = ["dep:hexf-parse", "dep:unicode-ident", "compact"]
+wgsl-in = ["dep:hexf-parse", "dep:unicode-ident"]
 wgsl-out = []
 
 ## Enables outputting to HLSL (Microsoft's High-Level Shader Language).
@@ -72,8 +72,6 @@ hlsl-out = []
 ## If you want to enable HLSL output it regardless of the target platform, use `naga/hlsl-out`.
 hlsl-out-if-target-windows = []
 
-compact = []
-
 ## Enables colored output through codespan-reporting and termcolor.
 termcolor = ["codespan-reporting/termcolor"]
 

naga/src/arena/handle_set.rs

@@ -77,6 +77,11 @@ impl<T> HandleSet<T> {
         }
     }
 
+    /// Add all of the handles that can be included in this set.
+    pub fn add_all(&mut self) {
+        self.members.get_mut().set_all();
+    }
+
     pub fn contains(&self, handle: Handle<T>) -> bool {
         self.members.contains(handle.index())
     }
@@ -85,6 +90,23 @@ impl<T> HandleSet<T> {
     pub fn iter(&self) -> impl '_ + Iterator<Item = Handle<T>> {
         self.members.iter().map(Handle::from_usize)
     }
+
+    /// Removes and returns the numerically largest handle in the set, or `None`
+    /// if the set is empty.
+    pub fn pop(&mut self) -> Option<Handle<T>> {
+        let members = core::mem::take(&mut self.members);
+        let mut vec = members.into_bit_vec();
+        let result = vec.iter_mut().enumerate().rev().find_map(|(i, mut bit)| {
+            if *bit {
+                *bit = false;
+                Some(i)
+            } else {
+                None
+            }
+        });
+        self.members = bit_set::BitSet::from_bit_vec(vec);
+        result.map(Handle::from_usize)
+    }
 }
 
 impl<T> Default for HandleSet<T> {

naga/src/arena/mod.rs

@@ -234,7 +234,6 @@ impl<T> Arena<T> {
         Ok(())
     }
 
-    #[cfg(feature = "compact")]
     pub(crate) fn retain_mut<P>(&mut self, mut predicate: P)
     where
         P: FnMut(Handle<T>, &mut T) -> bool,

naga/src/arena/unique_arena.rs

@@ -72,7 +72,6 @@ impl<T> UniqueArena<T> {
             .unwrap_or(&Span::default())
     }
 
-    #[cfg(feature = "compact")]
     pub(crate) fn drain_all(&mut self) -> UniqueArenaDrain<T> {
         UniqueArenaDrain {
             inner_elts: self.set.drain(..),
@@ -82,14 +81,12 @@ impl<T> UniqueArena<T> {
     }
 }
 
-#[cfg(feature = "compact")]
 pub struct UniqueArenaDrain<'a, T> {
     inner_elts: indexmap::set::Drain<'a, T>,
     inner_spans: alloc::vec::Drain<'a, Span>,
     index: Index,
 }
 
-#[cfg(feature = "compact")]
 impl<T> Iterator for UniqueArenaDrain<'_, T> {
     type Item = (Handle<T>, T, Span);
 

naga/src/back/pipeline_constants.rs

@@ -10,6 +10,8 @@ use thiserror::Error;
 use super::PipelineConstants;
 use crate::{
     arena::HandleVec,
+    compact::{compact, KeepUnused},
+    ir,
     proc::{ConstantEvaluator, ConstantEvaluatorError, Emitter},
     valid::{Capabilities, ModuleInfo, ValidationError, ValidationFlags, Validator},
     Arena, Block, Constant, Expression, Function, Handle, Literal, Module, Override, Range, Scalar,
@@ -55,6 +57,7 @@ pub enum PipelineConstantError {
 pub fn process_overrides<'a>(
     module: &'a Module,
     module_info: &'a ModuleInfo,
+    entry_point: Option<(ir::ShaderStage, &str)>,
     pipeline_constants: &PipelineConstants,
 ) -> Result<(Cow<'a, Module>, Cow<'a, ModuleInfo>), PipelineConstantError> {
     if module.overrides.is_empty() {
@@ -63,6 +66,16 @@ pub fn process_overrides<'a>(
 
     let mut module = module.clone();
 
+    // If an entry point was specified, compact the module to remove anything
+    // not reachable from that entry point. This is necessary because we may not
+    // have values for overrides that are not reachable from the entry point.
+    if let Some((ep_stage, ep_name)) = entry_point {
+        module
+            .entry_points
+            .retain(|ep| ep.stage == ep_stage && ep.name == ep_name);
+    }
+    compact(&mut module, KeepUnused::No);
+
     // A map from override handles to the handles of the constants
     // we've replaced them with.
     let mut override_map = HandleVec::with_capacity(module.overrides.len());

naga/src/compact/expressions.rs

@@ -11,6 +11,9 @@ pub struct ExpressionTracer<'tracer> {
     /// The used map for `types`.
     pub types_used: &'tracer mut HandleSet<crate::Type>,
 
+    /// The used map for global variables.
+    pub global_variables_used: &'tracer mut HandleSet<crate::GlobalVariable>,
+
     /// The used map for `constants`.
     pub constants_used: &'tracer mut HandleSet<crate::Constant>,
 
@@ -76,9 +79,7 @@ impl ExpressionTracer<'_> {
             // Expressions that do not contain handles that need to be traced.
             Ex::Literal(_)
             | Ex::FunctionArgument(_)
-            | Ex::GlobalVariable(_)
             | Ex::LocalVariable(_)
-            | Ex::CallResult(_)
             | Ex::SubgroupBallotResult
             | Ex::RayQueryProceedResult => {}
 
@@ -134,6 +135,9 @@ impl ExpressionTracer<'_> {
             } => {
                 self.expressions_used.insert(vector);
             }
+            Ex::GlobalVariable(handle) => {
+                self.global_variables_used.insert(handle);
+            }
             Ex::Load { pointer } => {
                 self.expressions_used.insert(pointer);
             }
@@ -233,6 +237,10 @@ impl ExpressionTracer<'_> {
             Ex::ArrayLength(expr) => {
                 self.expressions_used.insert(expr);
             }
+            // `CallResult` expressions do contain a function handle, but any used
+            // `CallResult` expression should have an associated `ir::Statement::Call`
+            // that we will trace.
+            Ex::CallResult(_) => {}
             Ex::AtomicResult { ty, comparison: _ }
             | Ex::WorkGroupUniformLoadResult { ty }
             | Ex::SubgroupOperationResult { ty } => {
@@ -267,9 +275,7 @@ impl ModuleMap {
             // Expressions that do not contain handles that need to be adjusted.
             Ex::Literal(_)
             | Ex::FunctionArgument(_)
-            | Ex::GlobalVariable(_)
             | Ex::LocalVariable(_)
-            | Ex::CallResult(_)
             | Ex::SubgroupBallotResult
             | Ex::RayQueryProceedResult => {}
 
@@ -306,6 +312,7 @@ impl ModuleMap {
                 ref mut vector,
                 pattern: _,
             } => adjust(vector),
+            Ex::GlobalVariable(ref mut handle) => self.globals.adjust(handle),
             Ex::Load { ref mut pointer } => adjust(pointer),
             Ex::ImageSample {
                 ref mut image,
@@ -392,6 +399,9 @@ impl ModuleMap {
                 kind: _,
                 convert: _,
             } => adjust(expr),
+            Ex::CallResult(ref mut function) => {
+                self.functions.adjust(function);
+            }
             Ex::AtomicResult {
                 ref mut ty,
                 comparison: _,

naga/src/compact/functions.rs

@@ -6,7 +6,10 @@ pub struct FunctionTracer<'a> {
     pub constants: &'a crate::Arena<crate::Constant>,
     pub overrides: &'a crate::Arena<crate::Override>,
 
+    pub functions_pending: &'a mut HandleSet<crate::Function>,
+    pub functions_used: &'a mut HandleSet<crate::Function>,
     pub types_used: &'a mut HandleSet<crate::Type>,
+    pub global_variables_used: &'a mut HandleSet<crate::GlobalVariable>,
     pub constants_used: &'a mut HandleSet<crate::Constant>,
     pub overrides_used: &'a mut HandleSet<crate::Override>,
     pub global_expressions_used: &'a mut HandleSet<crate::Expression>,
@@ -16,6 +19,13 @@ pub struct FunctionTracer<'a> {
 }
 
 impl FunctionTracer<'_> {
+    pub fn trace_call(&mut self, function: crate::Handle<crate::Function>) {
+        if !self.functions_used.contains(function) {
+            self.functions_used.insert(function);
+            self.functions_pending.insert(function);
+        }
+    }
+
     pub fn trace(&mut self) {
         for argument in self.function.arguments.iter() {
             self.types_used.insert(argument.ty);
@@ -53,6 +63,7 @@ impl FunctionTracer<'_> {
             expressions: &self.function.expressions,
 
             types_used: self.types_used,
+            global_variables_used: self.global_variables_used,
             constants_used: self.constants_used,
             overrides_used: self.overrides_used,
             expressions_used: &mut self.expressions_used,
@@ -105,6 +116,6 @@ impl FunctionMap {
         assert!(reuse.is_empty());
 
         // Adjust statements.
-        self.adjust_body(function);
+        self.adjust_body(function, &module_map.functions);
     }
 }