Skip to content

Commit

Permalink
[naga spv-in] Support atomics in fields of global structs
Browse files Browse the repository at this point in the history 
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 <[email protected]>
  • Loading branch information
@schell @jimblandy
schell and jimblandy committed Dec 17, 2024
1 parent 3160f81 commit 54dd512
Show file tree
Hide file tree
Showing 8 changed files with 267 additions and 109 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down
4 changes: 4 additions & 0 deletions naga/src/arena/handle_set.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,10 @@ impl<T> HandleSet<T> {
}
}

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<T>) -> Self {
let len = arena.len();
Expand Down
171 changes: 115 additions & 56 deletions naga/src/front/atomic_upgrade.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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")]
Expand Down Expand Up @@ -87,74 +85,131 @@ 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<GlobalVariable>,

/// 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<Handle<Type>, bit_set::BitSet>,
}

impl Upgrades {
pub fn insert_global(&mut self, global: Handle<GlobalVariable>) {
self.globals.insert(global);
}

pub fn insert_field(&mut self, struct_type: Handle<Type>, 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<Type>, Handle<Type>>,
}

impl UpgradeState<'_> {
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<Type>) -> Result<Handle<Type>, 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<Type>,
upgrades: &Upgrades,
) -> Result<Handle<Type>, 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) => {
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)?,
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 {
Expand All @@ -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<GlobalVariable>) -> 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(())
}
}
Expand All @@ -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<Item = Handle<GlobalVariable>>,
) -> 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(())
}
Expand Down
Loading

0 comments on commit 54dd512

Please sign in to comment.