[naga] Remove non-essential override references via compaction

.github/workflows/ci.yml

@@ -306,7 +306,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
 
       # Building for native platforms with standard tests.

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, false);
 
             // 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();

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:strum", "dep:unicode-ident", "compact"]
+wgsl-in = ["dep:hexf-parse", "dep:strum", "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"]
 
@@ -85,6 +83,7 @@ arbitrary = { workspace = true, features = ["derive"], optional = true }
 arrayvec.workspace = true
 bitflags.workspace = true
 bit-set.workspace = true
+bit-vec.workspace = true
 cfg-if.workspace = true
 codespan-reporting = { workspace = true }
 hashbrown.workspace = true

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,28 @@ 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()
+            .filter_map(|(i, mut bit)| {
+                if *bit {
+                    *bit = false;
+                    Some(i)
+                } else {
+                    None
+                }
+            })
+            .next();
+        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,
+    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, false);
+
     // 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);
     }
 }