[naga] Remove non-essential override references via compaction

Adds a mode to compaction that removes unused functions, global
variables, and named types and overrides. This mode is used
everywhere except the compaction at the end of lowering, where
it is important to preserve unused items for type checking and
other validation of the module.

Pruning all but the active entry point and then compacting makes
`process_overrides` tolerant of missing values for overrides that are
not used by the active entry point.

Fixes gfx-rs#5885

Author: andyleiserson

CHANGELOG.md

@@ -69,8 +69,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).
@@ -80,6 +79,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
 
@@ -107,6 +107,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
 
@@ -124,11 +128,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/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/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/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);
     }
 }

naga/src/compact/handle_set_map.rs

@@ -4,19 +4,43 @@ use crate::arena::{Arena, Handle, HandleSet, Range};
 
 type Index = crate::non_max_u32::NonMaxU32;
 
-/// A map from old handle indices to new, compressed handle indices.
-pub struct HandleMap<T> {
+/// A map keyed by handles.
+///
+/// In most cases, this is used to map from old handle indices to new,
+/// compressed handle indices.
+#[derive(Debug)]
+pub struct HandleMap<T, U = Index> {
     /// The indices assigned to handles in the compacted module.
     ///
     /// If `new_index[i]` is `Some(n)`, then `n` is the `Index` of the
     /// compacted `Handle` corresponding to the pre-compacted `Handle`
     /// whose index is `i`.
-    new_index: Vec<Option<Index>>,
+    new_index: Vec<Option<U>>,
 
     /// This type is indexed by values of type `T`.
     as_keys: core::marker::PhantomData<T>,
 }
 
+impl<T, U> HandleMap<T, U> {
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            new_index: Vec::with_capacity(capacity),
+            as_keys: core::marker::PhantomData,
+        }
+    }
+
+    pub fn get(&self, handle: Handle<T>) -> Option<&U> {
+        self.new_index.get(handle.index()).unwrap_or(&None).as_ref()
+    }
+
+    pub fn insert(&mut self, handle: Handle<T>, value: U) -> Option<U> {
+        if self.new_index.len() <= handle.index() {
+            self.new_index.resize_with(handle.index() + 1, || None);
+        }
+        core::mem::replace(&mut self.new_index[handle.index()], Some(value))
+    }
+}
+
 impl<T: 'static> HandleMap<T> {
     pub fn from_set(set: HandleSet<T>) -> Self {
         let mut next_index = Index::new(0).unwrap();