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 0000000000..62029146f7 Binary files /dev/null and b/naga/tests/in/spv/atomic_global_struct_field_vertex.spv differ diff --git a/naga/tests/in/spv/atomic_global_struct_field_vertex.spvasm b/naga/tests/in/spv/atomic_global_struct_field_vertex.spvasm new file mode 100644 index 0000000000..e70bd951c3 --- /dev/null +++ b/naga/tests/in/spv/atomic_global_struct_field_vertex.spvasm @@ -0,0 +1,53 @@ +; SPIR-V +; Version: 1.5 +; Generator: Google rspirv; 0 +; Bound: 44 +; Schema: 0 + OpCapability Shader + OpCapability VulkanMemoryModel + OpMemoryModel Logical Vulkan + OpEntryPoint Vertex %1 "global_field_vertex" %2 %gl_Position + OpMemberDecorate %_struct_9 0 Offset 0 + OpMemberDecorate %_struct_9 1 Offset 8 + OpMemberDecorate %_struct_9 2 Offset 16 + OpDecorate %_struct_10 Block + OpMemberDecorate %_struct_10 0 Offset 0 + OpDecorate %2 Binding 0 + OpDecorate %2 DescriptorSet 0 + OpDecorate %gl_Position BuiltIn Position + %uint = OpTypeInt 32 0 + %float = OpTypeFloat 32 + %v2float = OpTypeVector %float 2 + %_struct_9 = OpTypeStruct %uint %v2float %uint + %_struct_10 = OpTypeStruct %_struct_9 +%_ptr_StorageBuffer__struct_10 = OpTypePointer StorageBuffer %_struct_10 + %v4float = OpTypeVector %float 4 +%_ptr_Output_v4float = OpTypePointer Output %v4float + %void = OpTypeVoid + %18 = OpTypeFunction %void + %2 = OpVariable %_ptr_StorageBuffer__struct_10 StorageBuffer + %uint_0 = OpConstant %uint 0 +%_ptr_StorageBuffer_uint = OpTypePointer StorageBuffer %uint + %uint_2 = OpConstant %uint 2 + %uint_5 = OpConstant %uint 5 +%_ptr_StorageBuffer_v2float = OpTypePointer StorageBuffer %v2float + %uint_1 = OpConstant %uint 1 + %float_0 = OpConstant %float 0 +%gl_Position = OpVariable %_ptr_Output_v4float Output + %1 = OpFunction %void None %18 + %28 = OpLabel + %30 = OpInBoundsAccessChain %_ptr_StorageBuffer_uint %2 %uint_0 %uint_2 + %31 = OpInBoundsAccessChain %_ptr_StorageBuffer_uint %2 %uint_0 %uint_0 + %32 = OpLoad %uint %31 + %33 = OpAtomicIAdd %uint %30 %uint_5 %uint_0 %32 + %34 = OpConvertUToF %float %33 + %35 = OpInBoundsAccessChain %_ptr_StorageBuffer_v2float %2 %uint_0 %uint_1 + %36 = OpLoad %v2float %35 + %37 = OpCompositeExtract %float %36 0 + %38 = OpCompositeExtract %float %36 1 + %39 = OpFMul %float %34 %37 + %40 = OpFMul %float %34 %38 + %43 = OpCompositeConstruct %v4float %39 %40 %float_0 %34 + OpStore %gl_Position %43 + OpReturn + OpFunctionEnd diff --git a/naga/tests/out/wgsl/atomic_global_struct_field_vertex.wgsl b/naga/tests/out/wgsl/atomic_global_struct_field_vertex.wgsl new file mode 100644 index 0000000000..c3369e241b --- /dev/null +++ b/naga/tests/out/wgsl/atomic_global_struct_field_vertex.wgsl @@ -0,0 +1,29 @@ +struct type_5 { + member: u32, + member_1: vec2, + 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,