//! 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:
//!
//! - 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("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")]
GlobalVariableMissing,
#[error("atomic compare exchange requires a scalar base type")]
CompareExchangeNonScalarBaseType,
}
#[derive(Clone, Default)]
struct Padding(Arc<AtomicUsize>);
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()
}
}
#[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()
}
/// 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("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, upgrades)?,
space,
},
TypeInner::Array { base, size, stride } => TypeInner::Array {
base: self.upgrade_type(base, upgrades)?,
size,
stride,
},
TypeInner::Struct { ref members, span } => {
// 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");
};
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: new_members,
span,
}
}
TypeInner::BindingArray { base, size } => TypeInner::BindingArray {
base: self.upgrade_type(base, upgrades)?,
size,
},
_ => return Ok(ty),
};
// 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 {
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);
self.upgraded_types.insert(ty, new_handle);
Ok(new_handle)
}
fn upgrade_all(&mut self, upgrades: &Upgrades) -> Result<(), Error> {
for handle in upgrades.globals.iter() {
let padding = self.inc_padding();
let global = &self.module.global_variables[handle];
padding.trace("visiting global variable: ", handle);
padding.trace("var: ", global);
if global.init.is_some() {
return Err(Error::GlobalInitUnsupported);
}
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(())
}
}
impl Module {
/// Upgrade `global_var_handles` to have [`Atomic`] leaf types.
///
/// [`Atomic`]: TypeInner::Atomic
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(),
),
};
state.upgrade_all(upgrades)?;
Ok(())
}
}
