From 54dd512b422485d5a0ed47d4f9e9bceb1af1ff22 Mon Sep 17 00:00:00 2001 From: Schell Carl Scivally Date: Tue, 10 Dec 2024 11:07:41 +1300 Subject: [PATCH] [naga spv-in] Support atomics in fields of global structs Handle SPIR-V input that performs atomic operations on structs that have more than one field. Track which fields of which struct types are used by atomic operations on which global variables, and then give those global variables new types in which exactly those fields have had their `Scalar` leaf types changed to `Atomic`. Add snapshot tests. Co-authored-by: Jim Blandy --- CHANGELOG.md | 1 + naga/src/arena/handle_set.rs | 4 + naga/src/front/atomic_upgrade.rs | 171 ++++++++++++------ naga/src/front/spv/mod.rs | 117 ++++++------ .../spv/atomic_global_struct_field_vertex.spv | Bin 0 -> 820 bytes .../atomic_global_struct_field_vertex.spvasm | 53 ++++++ .../atomic_global_struct_field_vertex.wgsl | 29 +++ naga/tests/snapshots.rs | 1 + 8 files changed, 267 insertions(+), 109 deletions(-) create mode 100644 naga/tests/in/spv/atomic_global_struct_field_vertex.spv create mode 100644 naga/tests/in/spv/atomic_global_struct_field_vertex.spvasm create mode 100644 naga/tests/out/wgsl/atomic_global_struct_field_vertex.wgsl diff --git a/CHANGELOG.md b/CHANGELOG.md index 3429a9ff85..b89aac4aa7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -129,6 +129,7 @@ By @ErichDonGubler in [#6456](https://github.com/gfx-rs/wgpu/pull/6456), [#6148] #### Naga +- Support atomic operations on fields of global structs in the SPIR-V frontend. By @schell in [#6693](https://github.com/gfx-rs/wgpu/pull/6693). - Clean up tests for atomic operations support in SPIR-V frontend. By @schell in [#6692](https://github.com/gfx-rs/wgpu/pull/6692) - Fix an issue where `naga` CLI would incorrectly skip the first positional argument when `--stdin-file-path` was specified. By @ErichDonGubler in [#6480](https://github.com/gfx-rs/wgpu/pull/6480). - Fix textureNumLevels in the GLSL backend. By @magcius in [#6483](https://github.com/gfx-rs/wgpu/pull/6483). diff --git a/naga/src/arena/handle_set.rs b/naga/src/arena/handle_set.rs index ef2ded2ddb..f2ce058d12 100644 --- a/naga/src/arena/handle_set.rs +++ b/naga/src/arena/handle_set.rs @@ -25,6 +25,10 @@ impl HandleSet { } } + pub fn is_empty(&self) -> bool { + self.members.is_empty() + } + /// Return a new, empty `HandleSet`, sized to hold handles from `arena`. pub fn for_arena(arena: &impl ArenaType) -> Self { let len = arena.len(); diff --git a/naga/src/front/atomic_upgrade.rs b/naga/src/front/atomic_upgrade.rs index 171df33169..529883ae41 100644 --- a/naga/src/front/atomic_upgrade.rs +++ b/naga/src/front/atomic_upgrade.rs @@ -20,8 +20,6 @@ //! //! Future work: //! -//! - Atomics in structs are not implemented yet. -//! //! - The GLSL front end could use this transformation as well. //! //! [`Atomic`]: TypeInner::Atomic @@ -40,8 +38,8 @@ use crate::{GlobalVariable, Handle, Module, Type, TypeInner}; pub enum Error { #[error("encountered an unsupported expression")] Unsupported, - #[error("upgrading structs of more than one member is not yet implemented")] - MultiMemberStruct, + #[error("unexpected end of struct field access indices")] + UnexpectedEndOfIndices, #[error("encountered unsupported global initializer in an atomic variable")] GlobalInitUnsupported, #[error("expected to find a global variable")] @@ -87,9 +85,50 @@ impl Padding { } } +#[derive(Debug, Default)] +pub struct Upgrades { + /// Global variables that we've accessed using atomic operations. + /// + /// This includes globals with composite types (arrays, structs) where we've + /// only accessed some components (elements, fields) atomically. + globals: crate::arena::HandleSet, + + /// Struct fields that we've accessed using atomic operations. + /// + /// Each key refers to some [`Struct`] type, and each value is a set of + /// the indices of the fields in that struct that have been accessed + /// atomically. + /// + /// This includes fields with composite types (arrays, structs) + /// of which we've only accessed some components (elements, fields) + /// atomically. + /// + /// [`Struct`]: crate::TypeInner::Struct + fields: crate::FastHashMap, bit_set::BitSet>, +} + +impl Upgrades { + pub fn insert_global(&mut self, global: Handle) { + self.globals.insert(global); + } + + pub fn insert_field(&mut self, struct_type: Handle, field: usize) { + self.fields.entry(struct_type).or_default().insert(field); + } + + pub fn is_empty(&self) -> bool { + self.globals.is_empty() + } +} + struct UpgradeState<'a> { padding: Padding, module: &'a mut Module, + + /// A map from old types to their upgraded versions. + /// + /// This ensures we never try to rebuild a type more than once. + upgraded_types: crate::FastHashMap, Handle>, } impl UpgradeState<'_> { @@ -97,11 +136,38 @@ impl UpgradeState<'_> { self.padding.inc_padding() } - /// Upgrade the type, recursing until we reach the leaves. - /// At the leaves, replace scalars with atomic scalars. - fn upgrade_type(&mut self, ty: Handle) -> Result, Error> { + /// Get a type equivalent to `ty`, but with [`Scalar`] leaves upgraded to [`Atomic`] scalars. + /// + /// If such a type already exists in `self.module.types`, return its handle. + /// Otherwise, construct a new one and return that handle. + /// + /// If `ty` is a [`Pointer`], [`Array`], [`BindingArray`], recurse into the + /// type and upgrade its leaf types. + /// + /// If `ty` is a [`Struct`], recurse into it and upgrade only those fields + /// whose indices appear in `field_indices`. + /// + /// The existing type is not affected. + /// + /// [`Scalar`]: crate::TypeInner::Scalar + /// [`Atomic`]: crate::TypeInner::Atomic + /// [`Pointer`]: crate::TypeInner::Pointer + /// [`Array`]: crate::TypeInner::Array + /// [`Struct`]: crate::TypeInner::Struct + /// [`BindingArray`]: crate::TypeInner::BindingArray + fn upgrade_type( + &mut self, + ty: Handle, + upgrades: &Upgrades, + ) -> Result, Error> { let padding = self.inc_padding(); - padding.trace("upgrading type: ", ty); + padding.trace("visiting type: ", ty); + + // If we've already upgraded this type, return the handle we produced at + // the time. + if let Some(&new) = self.upgraded_types.get(&ty) { + return Ok(new); + } let inner = match self.module.types[ty].inner { TypeInner::Scalar(scalar) => { @@ -109,52 +175,41 @@ impl UpgradeState<'_> { TypeInner::Atomic(scalar) } TypeInner::Pointer { base, space } => TypeInner::Pointer { - base: self.upgrade_type(base)?, + base: self.upgrade_type(base, upgrades)?, space, }, TypeInner::Array { base, size, stride } => TypeInner::Array { - base: self.upgrade_type(base)?, + base: self.upgrade_type(base, upgrades)?, size, stride, }, TypeInner::Struct { ref members, span } => { - // In the future we should have to figure out which member needs - // upgrading, but for now we'll only cover the single-member - // case. - let &[crate::StructMember { - ref name, - ty, - ref binding, - offset, - }] = &members[..] - else { - return Err(Error::MultiMemberStruct); + // If no field or subfield of this struct was ever accessed + // atomically, no change is needed. We should never have arrived here. + let Some(fields) = upgrades.fields.get(&ty) else { + unreachable!("global or field incorrectly flagged as atomically accessed"); }; - // Take our own clones of these values now, so that - // `upgrade_type` can mutate the module. - let name = name.clone(); - let binding = binding.clone(); - let upgraded_member_type = self.upgrade_type(ty)?; + let mut new_members = members.clone(); + for field in fields { + new_members[field].ty = self.upgrade_type(new_members[field].ty, upgrades)?; + } + TypeInner::Struct { - members: vec![crate::StructMember { - name, - ty: upgraded_member_type, - binding, - offset, - }], + members: new_members, span, } } TypeInner::BindingArray { base, size } => TypeInner::BindingArray { - base: self.upgrade_type(base)?, + base: self.upgrade_type(base, upgrades)?, size, }, _ => return Ok(ty), }; - // Now that we've upgraded any subtypes, re-borrow a reference to our - // type and update its `inner`. + // At this point, we have a `TypeInner` that is the upgraded version of + // `ty`. Find a suitable `Type` for this, creating a new one if + // necessary, and return its handle. let r#type = &self.module.types[ty]; let span = self.module.types.get_span(ty); let new_type = Type { @@ -165,27 +220,32 @@ impl UpgradeState<'_> { padding.debug("from: ", r#type); padding.debug("to: ", &new_type); let new_handle = self.module.types.insert(new_type, span); + self.upgraded_types.insert(ty, new_handle); Ok(new_handle) } - fn upgrade_global_variable(&mut self, handle: Handle) -> Result<(), Error> { - let padding = self.inc_padding(); - padding.trace("upgrading global variable: ", handle); + fn upgrade_all(&mut self, upgrades: &Upgrades) -> Result<(), Error> { + for handle in upgrades.globals.iter() { + let padding = self.inc_padding(); - let var = &self.module.global_variables[handle]; + let global = &self.module.global_variables[handle]; + padding.trace("visiting global variable: ", handle); + padding.trace("var: ", global); - if var.init.is_some() { - return Err(Error::GlobalInitUnsupported); - } + if global.init.is_some() { + return Err(Error::GlobalInitUnsupported); + } - let var_ty = var.ty; - let new_ty = self.upgrade_type(var.ty)?; - if new_ty != var_ty { - padding.debug("upgrading global variable: ", handle); - padding.debug("from ty: ", var_ty); - padding.debug("to ty: ", new_ty); - self.module.global_variables[handle].ty = new_ty; + let var_ty = global.ty; + let new_ty = self.upgrade_type(var_ty, upgrades)?; + if new_ty != var_ty { + padding.debug("upgrading global variable: ", handle); + padding.debug("from ty: ", var_ty); + padding.debug("to ty: ", new_ty); + self.module.global_variables[handle].ty = new_ty; + } } + Ok(()) } } @@ -194,18 +254,17 @@ impl Module { /// Upgrade `global_var_handles` to have [`Atomic`] leaf types. /// /// [`Atomic`]: TypeInner::Atomic - pub(crate) fn upgrade_atomics( - &mut self, - global_var_handles: impl IntoIterator>, - ) -> Result<(), Error> { + pub(crate) fn upgrade_atomics(&mut self, upgrades: &Upgrades) -> Result<(), Error> { let mut state = UpgradeState { padding: Default::default(), module: self, + upgraded_types: crate::FastHashMap::with_capacity_and_hasher( + upgrades.fields.len(), + Default::default(), + ), }; - for handle in global_var_handles { - state.upgrade_global_variable(handle)?; - } + state.upgrade_all(upgrades)?; Ok(()) } diff --git a/naga/src/front/spv/mod.rs b/naga/src/front/spv/mod.rs index 587352973e..766ec0a8e7 100644 --- a/naga/src/front/spv/mod.rs +++ b/naga/src/front/spv/mod.rs @@ -36,7 +36,6 @@ mod null; use convert::*; pub use error::Error; use function::*; -use indexmap::IndexSet; use crate::{ arena::{Arena, Handle, UniqueArena}, @@ -47,6 +46,8 @@ use crate::{ use petgraph::graphmap::GraphMap; use std::{convert::TryInto, mem, num::NonZeroU32, path::PathBuf}; +use super::atomic_upgrade::Upgrades; + pub const SUPPORTED_CAPABILITIES: &[spirv::Capability] = &[ spirv::Capability::Shader, spirv::Capability::VulkanMemoryModel, @@ -557,45 +558,6 @@ struct BlockContext<'function> { parameter_sampling: &'function mut [image::SamplingFlags], } -impl BlockContext<'_> { - /// Descend into the expression with the given handle, locating a contained - /// global variable. - /// - /// If the expression doesn't actually refer to something in a global - /// variable, we can't upgrade its type in a way that Naga validation would - /// pass, so reject the input instead. - /// - /// This is used to track atomic upgrades. - fn get_contained_global_variable( - &self, - mut handle: Handle, - ) -> Result, Error> { - log::debug!("\t\tlocating global variable in {handle:?}"); - loop { - match self.expressions[handle] { - crate::Expression::Access { base, index: _ } => { - handle = base; - log::debug!("\t\t access {handle:?}"); - } - crate::Expression::AccessIndex { base, index: _ } => { - handle = base; - log::debug!("\t\t access index {handle:?}"); - } - crate::Expression::GlobalVariable(h) => { - log::debug!("\t\t found {h:?}"); - return Ok(h); - } - _ => { - break; - } - } - } - Err(Error::AtomicUpgradeError( - crate::front::atomic_upgrade::Error::GlobalVariableMissing, - )) - } -} - enum SignAnchor { Result, Operand, @@ -612,11 +574,12 @@ pub struct Frontend { future_member_decor: FastHashMap<(spirv::Word, MemberIndex), Decoration>, lookup_member: FastHashMap<(Handle, MemberIndex), LookupMember>, handle_sampling: FastHashMap, image::SamplingFlags>, - /// The set of all global variables accessed by [`Atomic`] statements we've + + /// A record of what is accessed by [`Atomic`] statements we've /// generated, so we can upgrade the types of their operands. /// /// [`Atomic`]: crate::Statement::Atomic - upgrade_atomics: IndexSet>, + upgrade_atomics: Upgrades, lookup_type: FastHashMap, lookup_void_type: Option, @@ -1481,8 +1444,7 @@ impl> Frontend { block.push(stmt, span); // Store any associated global variables so we can upgrade their types later - self.upgrade_atomics - .insert(ctx.get_contained_global_variable(p_lexp_handle)?); + self.record_atomic_access(ctx, p_lexp_handle)?; Ok(()) } @@ -4178,8 +4140,7 @@ impl> Frontend { ); // Store any associated global variables so we can upgrade their types later - self.upgrade_atomics - .insert(ctx.get_contained_global_variable(p_lexp_handle)?); + self.record_atomic_access(ctx, p_lexp_handle)?; } Op::AtomicStore => { inst.expect(5)?; @@ -4208,8 +4169,7 @@ impl> Frontend { emitter.start(ctx.expressions); // Store any associated global variables so we can upgrade their types later - self.upgrade_atomics - .insert(ctx.get_contained_global_variable(p_lexp_handle)?); + self.record_atomic_access(ctx, p_lexp_handle)?; } Op::AtomicIIncrement | Op::AtomicIDecrement => { inst.expect(6)?; @@ -4273,8 +4233,7 @@ impl> Frontend { block.push(stmt, span); // Store any associated global variables so we can upgrade their types later - self.upgrade_atomics - .insert(ctx.get_contained_global_variable(p_exp_h)?); + self.record_atomic_access(ctx, p_exp_h)?; } Op::AtomicCompareExchange => { inst.expect(9)?; @@ -4369,8 +4328,7 @@ impl> Frontend { block.push(stmt, span); // Store any associated global variables so we can upgrade their types later - self.upgrade_atomics - .insert(ctx.get_contained_global_variable(p_exp_h)?); + self.record_atomic_access(ctx, p_exp_h)?; } Op::AtomicExchange | Op::AtomicIAdd @@ -4685,7 +4643,7 @@ impl> Frontend { if !self.upgrade_atomics.is_empty() { log::info!("Upgrading atomic pointers..."); - module.upgrade_atomics(mem::take(&mut self.upgrade_atomics))?; + module.upgrade_atomics(&self.upgrade_atomics)?; } // Do entry point specific processing after all functions are parsed so that we can @@ -5980,6 +5938,59 @@ impl> Frontend { ); Ok(()) } + + /// Record an atomic access to some component of a global variable. + /// + /// Given `handle`, an expression referring to a scalar that has had an + /// atomic operation applied to it, descend into the expression, noting + /// which global variable it ultimately refers to, and which struct fields + /// of that global's value it accesses. + /// + /// Return the handle of the type of the expression. + /// + /// If the expression doesn't actually refer to something in a global + /// variable, we can't upgrade its type in a way that Naga validation would + /// pass, so reject the input instead. + fn record_atomic_access( + &mut self, + ctx: &BlockContext, + handle: Handle, + ) -> Result, Error> { + log::debug!("\t\tlocating global variable in {handle:?}"); + match ctx.expressions[handle] { + crate::Expression::Access { base, index } => { + log::debug!("\t\t access {handle:?} {index:?}"); + let ty = self.record_atomic_access(ctx, base)?; + let crate::TypeInner::Array { base, .. } = ctx.module.types[ty].inner else { + unreachable!("Atomic operations on Access expressions only work for arrays"); + }; + Ok(base) + } + crate::Expression::AccessIndex { base, index } => { + log::debug!("\t\t access index {handle:?} {index:?}"); + let ty = self.record_atomic_access(ctx, base)?; + match ctx.module.types[ty].inner { + crate::TypeInner::Struct { ref members, .. } => { + let index = index as usize; + self.upgrade_atomics.insert_field(ty, index); + Ok(members[index].ty) + } + crate::TypeInner::Array { base, .. } => { + Ok(base) + } + _ => unreachable!("Atomic operations on AccessIndex expressions only work for structs and arrays"), + } + } + crate::Expression::GlobalVariable(h) => { + log::debug!("\t\t found {h:?}"); + self.upgrade_atomics.insert_global(h); + Ok(ctx.module.global_variables[h].ty) + } + _ => Err(Error::AtomicUpgradeError( + crate::front::atomic_upgrade::Error::GlobalVariableMissing, + )), + } + } } fn make_index_literal( diff --git a/naga/tests/in/spv/atomic_global_struct_field_vertex.spv b/naga/tests/in/spv/atomic_global_struct_field_vertex.spv new file mode 100644 index 0000000000000000000000000000000000000000..62029146f7ffb8fa76f20eecd0fcad6ca09c5073 GIT binary patch literal 820 zcmYk4%}N7N429F5vHo_f+Ezuznbtq3YY{{jZt7B>peT+a3@9SS7w|#d`UGwS&v%=P zUPzOilbf8}X`?=DSyD5z)Q+y$GLMT~o3tFo~=Fa->LHu0~|*qVVW#B*Xkn57{{>yqbP z)f4BQ_tbyc^S-zh#pKM(TsQUL_{{aNPp_~)p2Ka+ILwb4(RrdS=%SW` z&TmcHk-xOLpe0kxY&4=H;=BKFxO3knR}A-Mt*}F(m>y`RE|9AlRcuRB!C4C*6x, + member_2: atomic, +} + +struct type_6 { + member: type_5, +} + +@group(0) @binding(0) +var global: type_6; +var global_1: vec4 = vec4(0f, 0f, 0f, 1f); + +fn function() { + let _e7 = global.member.member; + let _e8 = atomicAdd((&global.member.member_2), _e7); + let _e9 = f32(_e8); + let _e12 = global.member.member_1; + global_1 = vec4((_e9 * _e12.x), (_e9 * _e12.y), 0f, _e9); + return; +} + +@vertex +fn global_field_vertex() -> @builtin(position) vec4 { + function(); + let _e1 = global_1; + return _e1; +} diff --git a/naga/tests/snapshots.rs b/naga/tests/snapshots.rs index 72ce323585..e277ae7456 100644 --- a/naga/tests/snapshots.rs +++ b/naga/tests/snapshots.rs @@ -1077,6 +1077,7 @@ fn convert_spv_all() { convert_spv("atomic_compare_exchange", false, Targets::WGSL); convert_spv("atomic_i_decrement", false, Targets::WGSL); convert_spv("atomic_i_add_sub", false, Targets::WGSL); + convert_spv("atomic_global_struct_field_vertex", false, Targets::WGSL); convert_spv( "fetch_depth", false,