diff --git a/src/back/spv/layout.rs b/src/back/spv/layout.rs index 5708157ba5..492ec0f74a 100644 --- a/src/back/spv/layout.rs +++ b/src/back/spv/layout.rs @@ -25,6 +25,18 @@ impl PhysicalLayout { } } +impl super::recyclable::Recyclable for PhysicalLayout { + fn recycle(self) -> Self { + PhysicalLayout { + magic_number: self.magic_number, + version: self.version, + generator: self.generator, + instruction_schema: self.instruction_schema, + bound: 0, + } + } +} + impl LogicalLayout { pub(super) fn in_words(&self, sink: &mut impl Extend) { sink.extend(self.capabilities.iter().cloned()); @@ -41,6 +53,24 @@ impl LogicalLayout { } } +impl super::recyclable::Recyclable for LogicalLayout { + fn recycle(self) -> Self { + Self { + capabilities: self.capabilities.recycle(), + extensions: self.extensions.recycle(), + ext_inst_imports: self.ext_inst_imports.recycle(), + memory_model: self.memory_model.recycle(), + entry_points: self.entry_points.recycle(), + execution_modes: self.execution_modes.recycle(), + debugs: self.debugs.recycle(), + annotations: self.annotations.recycle(), + declarations: self.declarations.recycle(), + function_declarations: self.function_declarations.recycle(), + function_definitions: self.function_definitions.recycle(), + } + } +} + impl Instruction { pub(super) fn new(op: Op) -> Self { Instruction { diff --git a/src/back/spv/mod.rs b/src/back/spv/mod.rs index 0c3db44b2f..8445b22a8b 100644 --- a/src/back/spv/mod.rs +++ b/src/back/spv/mod.rs @@ -5,6 +5,7 @@ mod helpers; mod index; mod instructions; mod layout; +mod recyclable; mod writer; pub use spirv::Capability; @@ -15,6 +16,7 @@ use spirv::Word; use std::ops; use thiserror::Error; +#[derive(Clone)] struct PhysicalLayout { magic_number: Word, version: Word, @@ -239,6 +241,13 @@ impl ops::IndexMut> for CachedExpressions { id } } +impl recyclable::Recyclable for CachedExpressions { + fn recycle(self) -> Self { + CachedExpressions { + ids: self.ids.recycle(), + } + } +} struct GlobalVariable { /// Actual ID of the variable. diff --git a/src/back/spv/recyclable.rs b/src/back/spv/recyclable.rs new file mode 100644 index 0000000000..54ac912158 --- /dev/null +++ b/src/back/spv/recyclable.rs @@ -0,0 +1,74 @@ +//! Reusing collections' previous allocations. + +/// A value that can be reset to its initial state, retaining its current allocations. +/// +/// Naga attempts to lower the cost of SPIR-V generation by allowing clients to +/// reuse the same `Writer` for multiple Module translations. Reusing a `Writer` +/// means that the `Vec`s, `HashMap`s, and other heap-allocated structures the +/// `Writer` uses internally begin the translation with heap-allocated buffers +/// ready to use. +/// +/// But this approach introduces the risk of `Writer` state leaking from one +/// module to the next. When a developer adds fields to `Writer` or its internal +/// types, they must remember to reset their contents between modules. +/// +/// One trick to ensure that every field has been accounted for is to use Rust's +/// struct literal syntax to construct a new, reset value. If a developer adds a +/// field, but neglects to update the reset code, the compiler will complain +/// that a field is missing from the literal. This trait's `recycle` method +/// takes `self` by value, and returns `Self` by value, encouraging the use of +/// struct literal expressions in its implementation. +pub trait Recyclable { + /// Clear `self`, retaining its current memory allocations. + /// + /// Shrink the buffer if it's currently much larger than was actually used. + /// This prevents a module with exceptionally large allocations from causing + /// the `Writer` to retain more memory than it needs indefinitely. + fn recycle(self) -> Self; +} + +// Stock values for various collections. + +/// Maximum extra capacity that a recycled vector is allowed to have. If the +/// actual capacity is larger, we re-allocate the vector storage with lower +/// capacity. +const MAX_EXTRA_CAPACITY_PERCENT: usize = 200; + +/// Minimum extra capacity to keep when re-allocating the vector storage. +const MIN_EXTRA_CAPACITY_PERCENT: usize = 20; + +/// Minimum sensible length to consider for re-allocation. +const MIN_LENGTH: usize = 16; + +impl Recyclable for Vec { + fn recycle(mut self) -> Self { + let extra_capacity = (self.capacity() - self.len()) * 100 / self.len().max(MIN_LENGTH); + + if extra_capacity > MAX_EXTRA_CAPACITY_PERCENT { + //TODO: use `shrink_to` when it's stable + self = Vec::with_capacity(self.len() + self.len() * MIN_EXTRA_CAPACITY_PERCENT / 100); + } else { + self.clear(); + } + + self + } +} + +impl Recyclable for std::collections::HashMap { + fn recycle(mut self) -> Self { + let extra_capacity = (self.capacity() - self.len()) * 100 / self.len().max(MIN_LENGTH); + + if extra_capacity > MAX_EXTRA_CAPACITY_PERCENT { + //TODO: use `shrink_to` when it's stable + self = Self::with_capacity_and_hasher( + self.len() + self.len() * MIN_EXTRA_CAPACITY_PERCENT / 100, + self.hasher().clone(), + ); + } else { + self.clear(); + } + + self + } +} diff --git a/src/back/spv/writer.rs b/src/back/spv/writer.rs index 09f8fff8ee..5852da5656 100644 --- a/src/back/spv/writer.rs +++ b/src/back/spv/writer.rs @@ -118,9 +118,6 @@ impl Writer { return Err(Error::UnsupportedVersion(major, minor)); } let raw_version = ((major as u32) << 16) | ((minor as u32) << 8); - let mut id_gen = IdGenerator::default(); - let gl450_ext_inst_id = id_gen.next(); - let void_type = id_gen.next(); let (capabilities, forbidden_caps) = match options.capabilities { Some(ref caps) => (caps.clone(), None), @@ -132,6 +129,10 @@ impl Writer { } }; + let mut id_gen = IdGenerator::default(); + let gl450_ext_inst_id = id_gen.next(); + let void_type = id_gen.next(); + Ok(Writer { physical_layout: PhysicalLayout::new(raw_version), logical_layout: LogicalLayout::default(), @@ -156,6 +157,56 @@ impl Writer { }) } + /// Reset `Writer` to its initial state, retaining any allocations. + /// + /// Why not just implement `Recyclable` for `Writer`? By design, + /// `Recyclable::recycle` requires ownership of the value, not just + /// `&mut`; see the trait documentation. But we need to use this method + /// from functions like `Writer::write`, which only have `&mut Writer`. + /// Workarounds include unsafe code (`std::ptr::read`, then `write`, ugh) + /// or something like a `Default` impl that returns an oddly-initialized + /// `Writer`, which is worse. + fn reset(&mut self) { + use super::recyclable::Recyclable; + use std::mem::take; + + let mut id_gen = IdGenerator::default(); + let gl450_ext_inst_id = id_gen.next(); + let void_type = id_gen.next(); + + // Every field of the old writer that is not determined by the `Options` + // passed to `Writer::new` should be reset somehow. + let fresh = Writer { + // Copied from the old Writer: + flags: self.flags, + index_bounds_check_policy: self.index_bounds_check_policy, + capabilities: take(&mut self.capabilities), + forbidden_caps: take(&mut self.forbidden_caps), + + // Initialized afresh: + id_gen, + void_type, + gl450_ext_inst_id, + + // Recycled: + physical_layout: self.physical_layout.clone().recycle(), + logical_layout: take(&mut self.logical_layout).recycle(), + debugs: take(&mut self.debugs).recycle(), + annotations: take(&mut self.annotations).recycle(), + lookup_type: take(&mut self.lookup_type).recycle(), + lookup_function: take(&mut self.lookup_function).recycle(), + lookup_function_type: take(&mut self.lookup_function_type).recycle(), + lookup_function_call: take(&mut self.lookup_function_call).recycle(), + constant_ids: take(&mut self.constant_ids).recycle(), + cached_constants: take(&mut self.cached_constants).recycle(), + global_variables: take(&mut self.global_variables).recycle(), + cached: take(&mut self.cached).recycle(), + temp_list: take(&mut self.temp_list).recycle(), + }; + + *self = fresh; + } + fn check(&mut self, capabilities: &[spirv::Capability]) -> Result<(), Error> { if capabilities.is_empty() || capabilities @@ -2783,7 +2834,6 @@ impl Writer { .push(Instruction::source(spirv::SourceLanguage::GLSL, 450)); } - self.constant_ids.clear(); self.constant_ids.resize(ir_module.constants.len(), 0); // first, output all the scalar constants for (handle, constant) in ir_module.constants.iter() { @@ -2826,7 +2876,6 @@ impl Writer { debug_assert_eq!(self.constant_ids.iter().position(|&id| id == 0), None); // now write all globals - self.global_variables.clear(); for (_, var) in ir_module.global_variables.iter() { let (instruction, id) = self.write_global_variable(ir_module, var)?; instruction.to_words(&mut self.logical_layout.declarations); @@ -2884,9 +2933,7 @@ impl Writer { info: &ModuleInfo, words: &mut Vec, ) -> Result<(), Error> { - self.lookup_function.clear(); - self.lookup_function_type.clear(); - self.lookup_function_call.clear(); + self.reset(); self.write_logical_layout(ir_module, info)?; self.write_physical_layout();