From 967ad4dd7f86104636ca4bb27ab2d1ba4bc88181 Mon Sep 17 00:00:00 2001 From: Jim Blandy Date: Fri, 21 Jun 2024 21:26:30 -0700 Subject: [PATCH 1/2] [naga] Fix `cargo doc --document-private-items`. --- naga/src/arena.rs | 2 +- naga/src/back/mod.rs | 2 +- naga/src/back/msl/writer.rs | 2 ++ 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/naga/src/arena.rs b/naga/src/arena.rs index c442ad359d..7ca8ca9455 100644 --- a/naga/src/arena.rs +++ b/naga/src/arena.rs @@ -838,7 +838,7 @@ impl HandleVec { /// the end, like [`Vec::push`]. So the index of `handle` must equal /// [`self.len()`]. /// - /// [`HashMap]: std::collections::HashMap + /// [`HashMap`]: std::collections::HashMap /// [`self.len()`]: HandleVec::len pub(crate) fn insert(&mut self, handle: Handle, value: U) { assert_eq!(handle.index(), self.inner.len()); diff --git a/naga/src/back/mod.rs b/naga/src/back/mod.rs index 28a726155a..fb77b107c5 100644 --- a/naga/src/back/mod.rs +++ b/naga/src/back/mod.rs @@ -36,7 +36,7 @@ pub type NeedBakeExpressions = crate::FastHashSet); impl Display for ArraySizeMember { From 4e5066479f07bc20f74f3ffcf1be46756798feef Mon Sep 17 00:00:00 2001 From: Schell Carl Scivally Date: Wed, 5 Jun 2024 14:42:25 +1200 Subject: [PATCH 2/2] [naga spv-in] Adjust types of globals used by atomic instructions. To support atomic instructions in the SPIR-V front end, observe which global variables the input accesses using atomic instructions, and adjust their types from ordinary scalars to atomic values. See comments in `naga::front::atomic_upgrade`. --- CHANGELOG.md | 1 + naga/src/front/atomic_upgrade.rs | 214 ++++++++++++++++++ naga/src/front/mod.rs | 2 + naga/src/front/spv/error.rs | 5 +- naga/src/front/spv/mod.rs | 102 +++++---- naga/tests/in/spv/atomic_i_increment.spv | Bin 392 -> 828 bytes naga/tests/in/spv/atomic_i_increment.spvasm | 62 +++++ .../out/ir/atomic_i_increment.compact.ron | 202 ++++++++++++++++- naga/tests/out/ir/atomic_i_increment.ron | 201 ++++++++++++++-- 9 files changed, 724 insertions(+), 65 deletions(-) create mode 100644 naga/src/front/atomic_upgrade.rs create mode 100644 naga/tests/in/spv/atomic_i_increment.spvasm diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e617fe317..c6765c0e0b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -126,6 +126,7 @@ By @atlv24 in [#5383](https://github.com/gfx-rs/wgpu/pull/5383) #### Naga +- Added type upgrades to SPIR-V atomic support. Added related infrastructure. Tracking issue is [here](https://github.com/gfx-rs/wgpu/issues/4489). By @schell in [#5775](https://github.com/gfx-rs/wgpu/pull/5775). - Implement `WGSL`'s `unpack4xI8`,`unpack4xU8`,`pack4xI8` and `pack4xU8`. By @VlaDexa in [#5424](https://github.com/gfx-rs/wgpu/pull/5424) - Began work adding support for atomics to the SPIR-V frontend. Tracking issue is [here](https://github.com/gfx-rs/wgpu/issues/4489). By @schell in [#5702](https://github.com/gfx-rs/wgpu/pull/5702). diff --git a/naga/src/front/atomic_upgrade.rs b/naga/src/front/atomic_upgrade.rs new file mode 100644 index 0000000000..c59969ace9 --- /dev/null +++ b/naga/src/front/atomic_upgrade.rs @@ -0,0 +1,214 @@ +//! Upgrade the types of scalars observed to be accessed as atomics to [`Atomic`] types. +//! +//! In SPIR-V, atomic operations can be applied to any scalar value, but in Naga +//! IR atomic operations can only be applied to values of type [`Atomic`]. Naga +//! IR's restriction matches Metal Shading Language and WGSL, so we don't want +//! to relax that. Instead, when the SPIR-V front end observes a value being +//! accessed using atomic instructions, it promotes the value's type from +//! [`Scalar`] to [`Atomic`]. This module implements `Module::upgrade_atomics`, +//! the function that makes that change. +//! +//! Atomics can only appear in global variables in the [`Storage`] and +//! [`Workgroup`] address spaces. These variables can either have `Atomic` types +//! themselves, or be [`Array`]s of such, or be [`Struct`]s containing such. +//! So we only need to change the types of globals and struct fields. +//! +//! Naga IR [`Load`] expressions and [`Store`] statements can operate directly +//! on [`Atomic`] values, retrieving and depositing ordinary [`Scalar`] values, +//! so changing the types doesn't have much effect on the code that operates on +//! those values. +//! +//! Future work: +//! +//! - Atomics in structs are not implemented yet. +//! +//! - The GLSL front end could use this transformation as well. +//! +//! [`Atomic`]: TypeInner::Atomic +//! [`Scalar`]: TypeInner::Scalar +//! [`Storage`]: crate::AddressSpace::Storage +//! [`WorkGroup`]: crate::AddressSpace::WorkGroup +//! [`Array`]: TypeInner::Array +//! [`Struct`]: TypeInner::Struct +//! [`Load`]: crate::Expression::Load +//! [`Store`]: crate::Statement::Store +use std::sync::{atomic::AtomicUsize, Arc}; + +use crate::{GlobalVariable, Handle, Module, Type, TypeInner}; + +#[derive(Clone, Debug, thiserror::Error)] +pub enum Error { + #[error("encountered an unsupported expression")] + Unsupported, + #[error("upgrading structs of more than one member is not yet implemented")] + MultiMemberStruct, + #[error("encountered unsupported global initializer in an atomic variable")] + GlobalInitUnsupported, +} + +impl From for crate::front::spv::Error { + fn from(source: Error) -> Self { + crate::front::spv::Error::AtomicUpgradeError(source) + } +} + +#[derive(Clone, Default)] +struct Padding(Arc); + +impl std::fmt::Display for Padding { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + for _ in 0..self.0.load(std::sync::atomic::Ordering::Relaxed) { + f.write_str(" ")?; + } + Ok(()) + } +} + +impl Drop for Padding { + fn drop(&mut self) { + let _ = self.0.fetch_sub(1, std::sync::atomic::Ordering::Relaxed); + } +} + +impl Padding { + fn trace(&self, msg: impl std::fmt::Display, t: impl std::fmt::Debug) { + format!("{msg} {t:#?}") + .split('\n') + .for_each(|ln| log::trace!("{self}{ln}")); + } + + fn debug(&self, msg: impl std::fmt::Display, t: impl std::fmt::Debug) { + format!("{msg} {t:#?}") + .split('\n') + .for_each(|ln| log::debug!("{self}{ln}")); + } + + fn inc_padding(&self) -> Padding { + let _ = self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.clone() + } +} + +struct UpgradeState<'a> { + padding: Padding, + module: &'a mut Module, +} + +impl<'a> UpgradeState<'a> { + fn inc_padding(&self) -> Padding { + 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> { + let padding = self.inc_padding(); + padding.trace("upgrading type: ", ty); + + let inner = match self.module.types[ty].inner { + TypeInner::Scalar(scalar) => { + log::trace!("{padding}hit the scalar leaf, replacing with an atomic"); + TypeInner::Atomic(scalar) + } + TypeInner::Pointer { base, space } => TypeInner::Pointer { + base: self.upgrade_type(base)?, + space, + }, + TypeInner::Array { base, size, stride } => TypeInner::Array { + base: self.upgrade_type(base)?, + 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); + }; + + // 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)?; + TypeInner::Struct { + members: vec![crate::StructMember { + name, + ty: upgraded_member_type, + binding, + offset, + }], + span, + } + } + TypeInner::BindingArray { base, size } => TypeInner::BindingArray { + base: self.upgrade_type(base)?, + size, + }, + _ => return Ok(ty), + }; + + // Now that we've upgraded any subtypes, re-borrow a reference to our + // type and update its `inner`. + let r#type = &self.module.types[ty]; + let span = self.module.types.get_span(ty); + let new_type = Type { + name: r#type.name.clone(), + inner, + }; + padding.debug("ty: ", ty); + padding.debug("from: ", r#type); + padding.debug("to: ", &new_type); + let new_handle = self.module.types.insert(new_type, span); + Ok(new_handle) + } + + fn upgrade_global_variable(&mut self, handle: Handle) -> Result<(), Error> { + let padding = self.inc_padding(); + padding.trace("upgrading global variable: ", handle); + + let var = &self.module.global_variables[handle]; + + if var.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; + } + Ok(()) + } +} + +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> { + let mut state = UpgradeState { + padding: Default::default(), + module: self, + }; + + for handle in global_var_handles { + state.upgrade_global_variable(handle)?; + } + + Ok(()) + } +} diff --git a/naga/src/front/mod.rs b/naga/src/front/mod.rs index 8dc0416607..5e96103774 100644 --- a/naga/src/front/mod.rs +++ b/naga/src/front/mod.rs @@ -5,6 +5,8 @@ Frontend parsers that consume binary and text shaders and load them into [`Modul mod interpolator; mod type_gen; +#[cfg(feature = "spv-in")] +pub mod atomic_upgrade; #[cfg(feature = "glsl-in")] pub mod glsl; #[cfg(feature = "spv-in")] diff --git a/naga/src/front/spv/error.rs b/naga/src/front/spv/error.rs index af97ff5d97..42df9d807b 100644 --- a/naga/src/front/spv/error.rs +++ b/naga/src/front/spv/error.rs @@ -1,5 +1,5 @@ use super::ModuleState; -use crate::arena::Handle; +use crate::{arena::Handle, front::atomic_upgrade}; use codespan_reporting::diagnostic::Diagnostic; use codespan_reporting::files::SimpleFile; use codespan_reporting::term; @@ -134,6 +134,9 @@ pub enum Error { NonBindingArrayOfImageOrSamplers, #[error("naga only supports specialization constant IDs up to 65535 but was given {0}")] SpecIdTooHigh(u32), + + #[error("atomic upgrade error: {0}")] + AtomicUpgradeError(atomic_upgrade::Error), } impl Error { diff --git a/naga/src/front/spv/mod.rs b/naga/src/front/spv/mod.rs index 659667522f..d21811d1da 100644 --- a/naga/src/front/spv/mod.rs +++ b/naga/src/front/spv/mod.rs @@ -36,6 +36,7 @@ mod null; use convert::*; pub use error::Error; use function::*; +use indexmap::IndexSet; use crate::{ arena::{Arena, Handle, UniqueArena}, @@ -560,25 +561,44 @@ struct BlockContext<'function> { parameter_sampling: &'function mut [image::SamplingFlags], } +impl<'a> BlockContext<'a> { + /// Descend into the expression with the given handle, locating a contained + /// global variable. + /// + /// This is used to track atomic upgrades. + fn get_contained_global_variable( + &self, + mut handle: Handle, + ) -> Option> { + 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 Some(h); + } + _ => { + break; + } + } + } + None + } +} + enum SignAnchor { Result, Operand, } -enum AtomicOpInst { - AtomicIIncrement, -} - -#[allow(dead_code)] -struct AtomicOp { - instruction: AtomicOpInst, - result_type_id: spirv::Word, - result_id: spirv::Word, - pointer_id: spirv::Word, - scope_id: spirv::Word, - memory_semantics_id: spirv::Word, -} - pub struct Frontend { data: I, data_offset: usize, @@ -590,8 +610,12 @@ pub struct Frontend { future_member_decor: FastHashMap<(spirv::Word, MemberIndex), Decoration>, lookup_member: FastHashMap<(Handle, MemberIndex), LookupMember>, handle_sampling: FastHashMap, image::SamplingFlags>, - // Used to upgrade types used in atomic ops to atomic types, keyed by pointer id - lookup_atomic: FastHashMap, + /// The set of all global variables accessed by [`Atomic`] statements we've + /// generated, so we can upgrade the types of their operands. + /// + /// [`Atomic`]: crate::Statement::Atomic + upgrade_atomics: IndexSet>, + lookup_type: FastHashMap, lookup_void_type: Option, lookup_storage_buffer_types: FastHashMap, crate::StorageAccess>, @@ -647,7 +671,7 @@ impl> Frontend { future_member_decor: FastHashMap::default(), handle_sampling: FastHashMap::default(), lookup_member: FastHashMap::default(), - lookup_atomic: FastHashMap::default(), + upgrade_atomics: Default::default(), lookup_type: FastHashMap::default(), lookup_void_type: None, lookup_storage_buffer_types: FastHashMap::default(), @@ -3968,30 +3992,21 @@ impl> Frontend { let result_type_id = self.next()?; let result_id = self.next()?; let pointer_id = self.next()?; - let scope_id = self.next()?; - let memory_semantics_id = self.next()?; - // Store the op for a later pass where we "upgrade" the pointer type - let atomic = AtomicOp { - instruction: AtomicOpInst::AtomicIIncrement, - result_type_id, - result_id, - pointer_id, - scope_id, - memory_semantics_id, - }; - self.lookup_atomic.insert(pointer_id, atomic); + let _scope_id = self.next()?; + let _memory_semantics_id = self.next()?; log::trace!("\t\t\tlooking up expr {:?}", pointer_id); - let (p_lexp_handle, p_lexp_ty_id) = { let lexp = self.lookup_expression.lookup(pointer_id)?; let handle = get_expr_handle!(pointer_id, &lexp); (handle, lexp.type_id) }; + log::trace!("\t\t\tlooking up type {pointer_id:?}"); let p_ty = self.lookup_type.lookup(p_lexp_ty_id)?; let p_ty_base_id = p_ty.base_id.ok_or(Error::InvalidAccessType(p_lexp_ty_id))?; + log::trace!("\t\t\tlooking up base type {p_ty_base_id:?} of {p_ty:?}"); let p_base_ty = self.lookup_type.lookup(p_ty_base_id)?; @@ -4032,6 +4047,10 @@ impl> Frontend { result: Some(r_lexp_handle), }; block.push(stmt, span); + + // Store any associated global variables so we can upgrade their types later + self.upgrade_atomics + .extend(ctx.get_contained_global_variable(p_lexp_handle)); } _ => { return Err(Error::UnsupportedInstruction(self.state, inst.op)); @@ -4314,6 +4333,11 @@ impl> Frontend { }?; } + if !self.upgrade_atomics.is_empty() { + log::info!("Upgrading atomic pointers..."); + module.upgrade_atomics(std::mem::take(&mut self.upgrade_atomics))?; + } + // Do entry point specific processing after all functions are parsed so that we can // cull unused problematic builtins of gl_PerVertex. for (ep, fun_id) in mem::take(&mut self.deferred_entry_points) { @@ -5689,17 +5713,20 @@ mod test { #[cfg(all(feature = "wgsl-in", feature = "wgsl-out"))] #[test] fn atomic_i_inc() { - let _ = env_logger::builder() - .is_test(true) - .filter_level(log::LevelFilter::Trace) - .try_init(); + let _ = env_logger::builder().is_test(true).try_init(); let bytes = include_bytes!("../../../tests/in/spv/atomic_i_increment.spv"); let m = super::parse_u8_slice(bytes, &Default::default()).unwrap(); let mut validator = crate::valid::Validator::new( crate::valid::ValidationFlags::empty(), Default::default(), ); - let info = validator.validate(&m).unwrap(); + let info = match validator.validate(&m) { + Err(e) => { + log::error!("{}", e.emit_to_string("")); + return; + } + Ok(i) => i, + }; let wgsl = crate::back::wgsl::write_string(&m, &info, crate::back::wgsl::WriterFlags::empty()) .unwrap(); @@ -5709,15 +5736,14 @@ mod test { Ok(m) => m, Err(e) => { log::error!("{}", e.emit_to_string(&wgsl)); - // at this point we know atomics create invalid modules - // so simply bail - return; + panic!("invalid module"); } }; let mut validator = crate::valid::Validator::new(crate::valid::ValidationFlags::all(), Default::default()); if let Err(e) = validator.validate(&m) { log::error!("{}", e.emit_to_string(&wgsl)); + panic!("invalid generated wgsl"); } } } diff --git a/naga/tests/in/spv/atomic_i_increment.spv b/naga/tests/in/spv/atomic_i_increment.spv index 1b17c53065a7acbf9e846b656f26ab6b7eb83cb6..6d713ca45b83825a59be81bbe1878afbbe5f7bf1 100644 GIT binary patch literal 828 zcmZ9JNlF7z6h&W*#@^T@_Qs0MK%I+395~WlgdjmOP{Gs!+>R6X;7Bj>lLfBI>cv*Rv2KVMcszM9f7s%loazY>a`{hiUJ+-|3|NVcPAbcQ=Ep z?w~uky}Iw;^l#H(INipmE1=Sj8-zBX2p4}X%k7l*@zj)_2u4eqJ#P+Js z-kdv+%ibpSacT+XS%P^2Y%?_p>Nm(ufk#E}G`^>dXYkc}$c%lnSVHd{XifgIw~nul zgQ9mHzlHA|a~-4XH=i?i3Xie;a+R@nGOjS|+#+_&uGA~PNBtj6E`jIN)Wfr7T;8F~ zHSy&$_epgbbWVJQUH%o!x?Jvc4SPe|df@}v>;~wpcvkFd6W@P4p~3xZVe*-O=g%$B z^F_b4Z7jj$(N~i_#^ztxmYUj__i+zudw}1(Z!5pM{yQs2c^~iK&g8d0!5-mN-yWt8 Z-_iGx!~6PwQJwa2n3l;Iw_+Ap&?*`HjhL#qns0#-y3KS7n}-X<6rwqAF?wZiCE$ zWdV`h)9!^#aWr#frd%iKgudH)@n5~Q^_u>Gd1UL2<#jz5KKlVf+$BjvYRs7tRe8ky zF;z3#IkDD`zDK>C